From 206a9ce921367ca10091c79ad00e116f195850ff Mon Sep 17 00:00:00 2001 From: somew1nd Date: Thu, 13 Aug 2026 11:46:47 -0700 Subject: [PATCH] feat: initial release of ARTEMIS Android Agent & Automation Platform ARTEMIS is an autonomous multimodal Android agent and mobile UI automation platform engineered for physical devices and emulators. It combines state-of-the-art multimodal Vision-Language Models (VLMs) with hierarchical closed-loop planning, resilient self-healing, and deep Android subsystem integration. Key Features & Architectural Highlights: - Dual-Engine Orchestration: * Pro Engine: Multi-agent closed-loop (Planner -> Operator -> Validator/Failure Analyzer -> Summarizer) powered by LangGraph. * Flash Engine: Ultra-fast reactive loop for deterministic, low-latency UI execution (<35 steps). - Multimodal Perception & Grounding: * Hybrid visual screenshot grounding, real-time OCR extraction, and XML accessibility hierarchy parsing. * Dynamic-First, Coordinate-Fallback locating strategy for maximum resilience against UI drift. - Autonomous Self-Healing & Diagnostics: * Failure analyzer and safety net for popup dismissal, crash recovery, and auto-retry. * Diagnostic sub-agents for log analysis, video recording inspection, and object detection. - Unified CLI & Python SDK: * Rich CLI with commands: init, doctor, run, trace, batch, bench, server, and mcp. * Fluent Python SDK with Pydantic structured output models and package locking. - Multi-Provider LLM Architecture: * Support for Google Gemini, OpenAI, Anthropic Claude, Vertex AI, OpenRouter, xAI Grok, and local endpoints (Ollama/vLLM). * Centralized configuration in config/artemis.jsonc. - Native MCP & Observability: * Model Context Protocol (MCP) server endpoints for AI IDEs. * Full trace recording with before/after screenshots, action overlays, and execution replays. Change-Id: I430b93ee2b76b54cc5a3c6eeccd15d2cf29830a1 --- .dockerignore | 16 + .env.example | 19 + .gcloudignore | 17 + .gitignore | 109 + CONTRIBUTING.md | 103 + Dockerfile | 67 + LICENSE | 202 ++ Makefile | 56 + README.md | 399 +++ apps/README.md | 26 + apps/__init__.py | 1 + artemis/__init__.py | 41 + artemis/__main__.py | 18 + artemis/agents/__init__.py | 37 + artemis/agents/base.py | 146 + artemis/agents/checker/checker.json | 4 + artemis/agents/checker/checker.py | 460 ++++ artemis/agents/diagnoser/diagnoser.json | 22 + artemis/agents/diagnoser/diagnoser.py | 575 ++++ artemis/agents/explorer/constants.py | 72 + artemis/agents/explorer/explorer.json | 27 + artemis/agents/explorer/explorer.py | 2260 ++++++++++++++++ artemis/agents/flash/context_compressor.py | 118 + artemis/agents/flash/flash_runner.md | 41 + artemis/agents/flash/flash_summarizer.md | 41 + artemis/agents/flash/runner.py | 524 ++++ artemis/agents/flash/summarizer.py | 226 ++ .../history_analyzer/history_analyzer.md | 8 + .../history_analyzer/history_analyzer.py | 250 ++ artemis/agents/hopper/hopper.md | 15 + artemis/agents/hopper/hopper.py | 93 + .../agents/image_processor/image_processor.md | 74 + .../agents/image_processor/image_processor.py | 322 +++ artemis/agents/log_analyzer/log_analyzer.md | 15 + artemis/agents/log_analyzer/log_analyzer.py | 431 +++ .../log_analyzer/log_reader_sub_agent.md | 8 + .../agents/log_analyzer/output_analyzer.md | 11 + .../agents/log_analyzer/output_analyzer.py | 162 ++ .../object_detector/object_detector.json | 9 + .../agents/object_detector/object_detector.py | 228 ++ artemis/agents/operator/loop_detector.py | 42 + artemis/agents/operator/operator.json | 4 + artemis/agents/operator/operator.py | 1481 +++++++++++ artemis/agents/operator/operator_agent.py | 72 + artemis/agents/operator/prompt_builder.py | 59 + artemis/agents/operator/prompts.py | 376 +++ artemis/agents/outputter/outputter.md | 28 + artemis/agents/outputter/outputter.py | 317 +++ artemis/agents/outputter/tools.py | 170 ++ artemis/agents/planner/planner.json | 30 + artemis/agents/planner/planner.py | 422 +++ artemis/agents/planner/planner_agent.py | 53 + artemis/agents/planner/types.py | 85 + artemis/agents/planner/utils.py | 89 + artemis/agents/summarizer/summarizer.md | 31 + artemis/agents/summarizer/summarizer.py | 231 ++ artemis/agents/summarizer/summarizer_agent.py | 53 + artemis/agents/validator/failure_analyzer.md | 25 + artemis/agents/validator/failure_analyzer.py | 607 +++++ .../validator/failure_analyzer_agent.py | 69 + artemis/agents/validator/pixel_safety_net.md | 11 + .../pixel_target_disappeared_analyzer.md | 65 + .../agents/validator/recovery_strategies.py | 177 ++ artemis/agents/validator/safety_net.py | 48 + .../validator/target_disappeared_analyzer.md | 61 + artemis/agents/validator/tool_declarations.py | 1180 +++++++++ artemis/agents/validator/validator.py | 1615 ++++++++++++ artemis/agents/video_analyzer/__init__.py | 15 + .../video_analyzer/conflict_resolution.py | 137 + .../agents/video_analyzer/video_analyzer.md | 51 + .../agents/video_analyzer/video_analyzer.py | 2347 +++++++++++++++++ .../agents/video_analyzer/video_sub_agent.md | 15 + .../video_analyzer/video_sub_agent_audio.md | 12 + artemis/clients/adb_tunnel.py | 327 +++ artemis/clients/ui_automator_client.py | 416 +++ artemis/config/__init__.py | 253 ++ artemis/config/agent.py | 490 ++++ artemis/config/constants.py | 142 + artemis/config/llm.py | 364 +++ artemis/config/output.py | 159 ++ artemis/config/paths.py | 170 ++ artemis/config/runtime.py | 142 + artemis/config/settings.py | 199 ++ artemis/constants.py | 50 + artemis/context.py | 219 ++ artemis/controllers/__init__.py | 41 + artemis/controllers/android_controller.py | 912 +++++++ artemis/controllers/controller_factory.py | 29 + artemis/controllers/device_controller.py | 245 ++ .../platform_specific_commands_controller.py | 137 + artemis/controllers/types.py | 37 + artemis/controllers/unified_controller.py | 796 ++++++ artemis/core/__init__.py | 51 + artemis/core/context.py | 59 + artemis/core/events.py | 77 + artemis/core/exceptions.py | 57 + artemis/core/registry.py | 46 + artemis/core/state.py | 70 + artemis/core/tool_declaration.py | 46 + artemis/data_engine/context_vars.py | 24 + artemis/data_engine/engine.py | 1138 ++++++++ artemis/data_engine/models.py | 102 + artemis/data_engine/storage.py | 1007 +++++++ artemis/data_engine/trace.py | 675 +++++ artemis/drivers/__init__.py | 48 + artemis/drivers/android/adb_driver.py | 481 ++++ artemis/drivers/android/hierarchy.py | 64 + artemis/drivers/android/input_ime.py | 119 + artemis/drivers/android/recorder.py | 57 + artemis/drivers/base.py | 271 ++ artemis/drivers/cloud/cloud_driver.py | 117 + artemis/drivers/factory.py | 78 + artemis/drivers/mock/mock_driver.py | 198 ++ artemis/drivers/types.py | 109 + artemis/engine/__init__.py | 27 + artemis/engine/base_runner.py | 36 + artemis/engine/graph_runner.py | 89 + artemis/engine/pipeline.py | 91 + artemis/engine/reactive_runner.py | 111 + artemis/graph/graph.py | 823 ++++++ artemis/graph/perception.py | 187 ++ artemis/graph/state.py | 134 + artemis/interfaces/cli/__init__.py | 19 + artemis/interfaces/cli/commands/__init__.py | 15 + artemis/interfaces/cli/commands/batch.py | 160 ++ artemis/interfaces/cli/commands/bench.py | 157 ++ artemis/interfaces/cli/commands/doctor.py | 250 ++ artemis/interfaces/cli/commands/init.py | 150 ++ artemis/interfaces/cli/commands/mcp.py | 42 + artemis/interfaces/cli/commands/run.py | 335 +++ artemis/interfaces/cli/commands/server.py | 137 + artemis/interfaces/cli/commands/trace.py | 111 + artemis/interfaces/cli/main.py | 99 + artemis/interfaces/sdk/__init__.py | 24 + artemis/interfaces/sdk/client.py | 147 ++ artemis/interfaces/sdk/task.py | 58 + artemis/llm/__init__.py | 29 + artemis/llm/client_factory.py | 25 + artemis/llm/parser.py | 32 + artemis/llm/router.py | 433 +++ artemis/llm/schemas.py | 36 + artemis/main.py | 41 + artemis/mcp/adb_server.py | 465 ++++ artemis/mcp/xml_search_server.py | 280 ++ artemis/sdk/__init__.py | 26 + artemis/sdk/agent.py | 1028 ++++++++ artemis/sdk/builders/__init__.py | 29 + artemis/sdk/builders/agent_config_builder.py | 528 ++++ artemis/sdk/builders/index.py | 29 + artemis/sdk/builders/task_request_builder.py | 256 ++ artemis/sdk/constants.py | 15 + artemis/sdk/examples/__init__.py | 15 + artemis/sdk/examples/app_lock_messaging.py | 87 + .../sdk/examples/simple_photo_organizer.py | 111 + .../examples/smart_notification_assistant.py | 235 ++ .../examples/video_transcription_example.py | 109 + artemis/sdk/types/__init__.py | 61 + artemis/sdk/types/agent.py | 145 + artemis/sdk/types/exceptions.py | 105 + artemis/sdk/types/platform.py | 210 ++ artemis/sdk/types/task.py | 246 ++ artemis/sdk/utils.py | 28 + artemis/services/llm.py | 635 +++++ artemis/telemetry/__init__.py | 35 + artemis/telemetry/models.py | 64 + artemis/telemetry/replay.py | 34 + artemis/telemetry/storage/json_storage.py | 47 + artemis/telemetry/tracer.py | 85 + artemis/tools/actions/__init__.py | 45 + artemis/tools/actions/device_actions.py | 237 ++ artemis/tools/base.py | 274 ++ artemis/tools/command_tool.py | 836 ++++++ artemis/tools/committee_prompts.json | 5 + artemis/tools/committee_tool.py | 427 +++ artemis/tools/diagnoser_submit_answer_tool.py | 88 + artemis/tools/diagnostic_tool.py | 120 + artemis/tools/explorer_tool.py | 404 +++ artemis/tools/image_processor_tool.py | 125 + artemis/tools/index.py | 52 + artemis/tools/log_tool.py | 108 + artemis/tools/memory/notes.py | 66 + artemis/tools/mobile/exec_tools.py | 730 +++++ artemis/tools/mobile/launch_app.py | 215 ++ artemis/tools/mobile/log_utils.py | 138 + artemis/tools/mobile/ocr.py | 210 ++ artemis/tools/mobile/read_hierarchy.py | 106 + artemis/tools/mobile/read_logs.py | 116 + artemis/tools/mobile/search_logs.py | 237 ++ artemis/tools/object_detection_tool.py | 246 ++ artemis/tools/perception/object_detect.py | 43 + artemis/tools/perception/ocr.py | 46 + artemis/tools/perception/xml_search.py | 45 + artemis/tools/scratchpad.py | 903 +++++++ artemis/tools/system/adb_shell.py | 31 + artemis/tools/tool_wrapper.py | 179 ++ artemis/tools/types.py | 72 + artemis/tools/utils.py | 389 +++ artemis/tools/video_tool.py | 253 ++ artemis/tools/wait_tool.py | 98 + artemis/utils/app_launch_utils.py | 255 ++ artemis/utils/cli_helpers.py | 38 + artemis/utils/cli_selection.py | 154 ++ artemis/utils/conversations.py | 45 + artemis/utils/coordinates.py | 344 +++ artemis/utils/cv_canvas.py | 189 ++ artemis/utils/decorators.py | 138 + artemis/utils/file.py | 55 + artemis/utils/image_diff.py | 335 +++ artemis/utils/logger.py | 212 ++ artemis/utils/media.py | 207 ++ artemis/utils/notes.py | 441 ++++ artemis/utils/ocr_api.py | 90 + artemis/utils/ocr_xml_fusion.py | 318 +++ artemis/utils/python_executor.py | 105 + artemis/utils/recorder.py | 60 + artemis/utils/requests_utils.py | 51 + artemis/utils/shell_utils.py | 34 + artemis/utils/task_tree.py | 1077 ++++++++ artemis/utils/text.py | 37 + artemis/utils/time.py | 20 + artemis/utils/ui_filter.py | 646 +++++ artemis/utils/ui_hierarchy.py | 166 ++ artemis/utils/verification.py | 62 + artemis/utils/video.py | 536 ++++ artemis/utils/visualization.py | 693 +++++ config/artemis.jsonc | 217 ++ pyproject.toml | 218 ++ pyrightconfig.json | 3 + setup.py | 104 + 229 files changed, 49653 insertions(+) create mode 100644 .dockerignore create mode 100644 .env.example create mode 100644 .gcloudignore create mode 100644 .gitignore create mode 100644 CONTRIBUTING.md create mode 100644 Dockerfile create mode 100644 LICENSE create mode 100644 Makefile create mode 100644 README.md create mode 100644 apps/README.md create mode 100644 apps/__init__.py create mode 100644 artemis/__init__.py create mode 100644 artemis/__main__.py create mode 100644 artemis/agents/__init__.py create mode 100644 artemis/agents/base.py create mode 100644 artemis/agents/checker/checker.json create mode 100644 artemis/agents/checker/checker.py create mode 100644 artemis/agents/diagnoser/diagnoser.json create mode 100644 artemis/agents/diagnoser/diagnoser.py create mode 100644 artemis/agents/explorer/constants.py create mode 100644 artemis/agents/explorer/explorer.json create mode 100644 artemis/agents/explorer/explorer.py create mode 100644 artemis/agents/flash/context_compressor.py create mode 100644 artemis/agents/flash/flash_runner.md create mode 100644 artemis/agents/flash/flash_summarizer.md create mode 100644 artemis/agents/flash/runner.py create mode 100644 artemis/agents/flash/summarizer.py create mode 100644 artemis/agents/history_analyzer/history_analyzer.md create mode 100644 artemis/agents/history_analyzer/history_analyzer.py create mode 100644 artemis/agents/hopper/hopper.md create mode 100644 artemis/agents/hopper/hopper.py create mode 100644 artemis/agents/image_processor/image_processor.md create mode 100644 artemis/agents/image_processor/image_processor.py create mode 100644 artemis/agents/log_analyzer/log_analyzer.md create mode 100644 artemis/agents/log_analyzer/log_analyzer.py create mode 100644 artemis/agents/log_analyzer/log_reader_sub_agent.md create mode 100644 artemis/agents/log_analyzer/output_analyzer.md create mode 100644 artemis/agents/log_analyzer/output_analyzer.py create mode 100644 artemis/agents/object_detector/object_detector.json create mode 100644 artemis/agents/object_detector/object_detector.py create mode 100644 artemis/agents/operator/loop_detector.py create mode 100644 artemis/agents/operator/operator.json create mode 100644 artemis/agents/operator/operator.py create mode 100644 artemis/agents/operator/operator_agent.py create mode 100644 artemis/agents/operator/prompt_builder.py create mode 100644 artemis/agents/operator/prompts.py create mode 100644 artemis/agents/outputter/outputter.md create mode 100644 artemis/agents/outputter/outputter.py create mode 100644 artemis/agents/outputter/tools.py create mode 100644 artemis/agents/planner/planner.json create mode 100644 artemis/agents/planner/planner.py create mode 100644 artemis/agents/planner/planner_agent.py create mode 100644 artemis/agents/planner/types.py create mode 100644 artemis/agents/planner/utils.py create mode 100644 artemis/agents/summarizer/summarizer.md create mode 100644 artemis/agents/summarizer/summarizer.py create mode 100644 artemis/agents/summarizer/summarizer_agent.py create mode 100644 artemis/agents/validator/failure_analyzer.md create mode 100644 artemis/agents/validator/failure_analyzer.py create mode 100644 artemis/agents/validator/failure_analyzer_agent.py create mode 100644 artemis/agents/validator/pixel_safety_net.md create mode 100644 artemis/agents/validator/pixel_target_disappeared_analyzer.md create mode 100644 artemis/agents/validator/recovery_strategies.py create mode 100644 artemis/agents/validator/safety_net.py create mode 100644 artemis/agents/validator/target_disappeared_analyzer.md create mode 100644 artemis/agents/validator/tool_declarations.py create mode 100644 artemis/agents/validator/validator.py create mode 100644 artemis/agents/video_analyzer/__init__.py create mode 100644 artemis/agents/video_analyzer/conflict_resolution.py create mode 100644 artemis/agents/video_analyzer/video_analyzer.md create mode 100644 artemis/agents/video_analyzer/video_analyzer.py create mode 100644 artemis/agents/video_analyzer/video_sub_agent.md create mode 100644 artemis/agents/video_analyzer/video_sub_agent_audio.md create mode 100644 artemis/clients/adb_tunnel.py create mode 100644 artemis/clients/ui_automator_client.py create mode 100644 artemis/config/__init__.py create mode 100644 artemis/config/agent.py create mode 100644 artemis/config/constants.py create mode 100644 artemis/config/llm.py create mode 100644 artemis/config/output.py create mode 100644 artemis/config/paths.py create mode 100644 artemis/config/runtime.py create mode 100644 artemis/config/settings.py create mode 100644 artemis/constants.py create mode 100644 artemis/context.py create mode 100644 artemis/controllers/__init__.py create mode 100644 artemis/controllers/android_controller.py create mode 100644 artemis/controllers/controller_factory.py create mode 100644 artemis/controllers/device_controller.py create mode 100644 artemis/controllers/platform_specific_commands_controller.py create mode 100644 artemis/controllers/types.py create mode 100644 artemis/controllers/unified_controller.py create mode 100644 artemis/core/__init__.py create mode 100644 artemis/core/context.py create mode 100644 artemis/core/events.py create mode 100644 artemis/core/exceptions.py create mode 100644 artemis/core/registry.py create mode 100644 artemis/core/state.py create mode 100644 artemis/core/tool_declaration.py create mode 100644 artemis/data_engine/context_vars.py create mode 100644 artemis/data_engine/engine.py create mode 100644 artemis/data_engine/models.py create mode 100644 artemis/data_engine/storage.py create mode 100644 artemis/data_engine/trace.py create mode 100644 artemis/drivers/__init__.py create mode 100644 artemis/drivers/android/adb_driver.py create mode 100644 artemis/drivers/android/hierarchy.py create mode 100644 artemis/drivers/android/input_ime.py create mode 100644 artemis/drivers/android/recorder.py create mode 100644 artemis/drivers/base.py create mode 100644 artemis/drivers/cloud/cloud_driver.py create mode 100644 artemis/drivers/factory.py create mode 100644 artemis/drivers/mock/mock_driver.py create mode 100644 artemis/drivers/types.py create mode 100644 artemis/engine/__init__.py create mode 100644 artemis/engine/base_runner.py create mode 100644 artemis/engine/graph_runner.py create mode 100644 artemis/engine/pipeline.py create mode 100644 artemis/engine/reactive_runner.py create mode 100644 artemis/graph/graph.py create mode 100644 artemis/graph/perception.py create mode 100644 artemis/graph/state.py create mode 100644 artemis/interfaces/cli/__init__.py create mode 100644 artemis/interfaces/cli/commands/__init__.py create mode 100644 artemis/interfaces/cli/commands/batch.py create mode 100644 artemis/interfaces/cli/commands/bench.py create mode 100644 artemis/interfaces/cli/commands/doctor.py create mode 100644 artemis/interfaces/cli/commands/init.py create mode 100644 artemis/interfaces/cli/commands/mcp.py create mode 100644 artemis/interfaces/cli/commands/run.py create mode 100644 artemis/interfaces/cli/commands/server.py create mode 100644 artemis/interfaces/cli/commands/trace.py create mode 100644 artemis/interfaces/cli/main.py create mode 100644 artemis/interfaces/sdk/__init__.py create mode 100644 artemis/interfaces/sdk/client.py create mode 100644 artemis/interfaces/sdk/task.py create mode 100644 artemis/llm/__init__.py create mode 100644 artemis/llm/client_factory.py create mode 100644 artemis/llm/parser.py create mode 100644 artemis/llm/router.py create mode 100644 artemis/llm/schemas.py create mode 100644 artemis/main.py create mode 100644 artemis/mcp/adb_server.py create mode 100644 artemis/mcp/xml_search_server.py create mode 100644 artemis/sdk/__init__.py create mode 100644 artemis/sdk/agent.py create mode 100644 artemis/sdk/builders/__init__.py create mode 100644 artemis/sdk/builders/agent_config_builder.py create mode 100644 artemis/sdk/builders/index.py create mode 100644 artemis/sdk/builders/task_request_builder.py create mode 100644 artemis/sdk/constants.py create mode 100644 artemis/sdk/examples/__init__.py create mode 100644 artemis/sdk/examples/app_lock_messaging.py create mode 100644 artemis/sdk/examples/simple_photo_organizer.py create mode 100644 artemis/sdk/examples/smart_notification_assistant.py create mode 100644 artemis/sdk/examples/video_transcription_example.py create mode 100644 artemis/sdk/types/__init__.py create mode 100644 artemis/sdk/types/agent.py create mode 100644 artemis/sdk/types/exceptions.py create mode 100644 artemis/sdk/types/platform.py create mode 100644 artemis/sdk/types/task.py create mode 100644 artemis/sdk/utils.py create mode 100644 artemis/services/llm.py create mode 100644 artemis/telemetry/__init__.py create mode 100644 artemis/telemetry/models.py create mode 100644 artemis/telemetry/replay.py create mode 100644 artemis/telemetry/storage/json_storage.py create mode 100644 artemis/telemetry/tracer.py create mode 100644 artemis/tools/actions/__init__.py create mode 100644 artemis/tools/actions/device_actions.py create mode 100644 artemis/tools/base.py create mode 100644 artemis/tools/command_tool.py create mode 100644 artemis/tools/committee_prompts.json create mode 100644 artemis/tools/committee_tool.py create mode 100644 artemis/tools/diagnoser_submit_answer_tool.py create mode 100644 artemis/tools/diagnostic_tool.py create mode 100644 artemis/tools/explorer_tool.py create mode 100644 artemis/tools/image_processor_tool.py create mode 100644 artemis/tools/index.py create mode 100644 artemis/tools/log_tool.py create mode 100644 artemis/tools/memory/notes.py create mode 100644 artemis/tools/mobile/exec_tools.py create mode 100644 artemis/tools/mobile/launch_app.py create mode 100644 artemis/tools/mobile/log_utils.py create mode 100644 artemis/tools/mobile/ocr.py create mode 100644 artemis/tools/mobile/read_hierarchy.py create mode 100644 artemis/tools/mobile/read_logs.py create mode 100644 artemis/tools/mobile/search_logs.py create mode 100644 artemis/tools/object_detection_tool.py create mode 100644 artemis/tools/perception/object_detect.py create mode 100644 artemis/tools/perception/ocr.py create mode 100644 artemis/tools/perception/xml_search.py create mode 100644 artemis/tools/scratchpad.py create mode 100644 artemis/tools/system/adb_shell.py create mode 100644 artemis/tools/tool_wrapper.py create mode 100644 artemis/tools/types.py create mode 100644 artemis/tools/utils.py create mode 100644 artemis/tools/video_tool.py create mode 100644 artemis/tools/wait_tool.py create mode 100644 artemis/utils/app_launch_utils.py create mode 100644 artemis/utils/cli_helpers.py create mode 100644 artemis/utils/cli_selection.py create mode 100644 artemis/utils/conversations.py create mode 100644 artemis/utils/coordinates.py create mode 100644 artemis/utils/cv_canvas.py create mode 100644 artemis/utils/decorators.py create mode 100644 artemis/utils/file.py create mode 100644 artemis/utils/image_diff.py create mode 100644 artemis/utils/logger.py create mode 100644 artemis/utils/media.py create mode 100644 artemis/utils/notes.py create mode 100644 artemis/utils/ocr_api.py create mode 100644 artemis/utils/ocr_xml_fusion.py create mode 100644 artemis/utils/python_executor.py create mode 100644 artemis/utils/recorder.py create mode 100644 artemis/utils/requests_utils.py create mode 100644 artemis/utils/shell_utils.py create mode 100644 artemis/utils/task_tree.py create mode 100644 artemis/utils/text.py create mode 100644 artemis/utils/time.py create mode 100644 artemis/utils/ui_filter.py create mode 100644 artemis/utils/ui_hierarchy.py create mode 100644 artemis/utils/verification.py create mode 100644 artemis/utils/video.py create mode 100644 artemis/utils/visualization.py create mode 100644 config/artemis.jsonc create mode 100644 pyproject.toml create mode 100644 pyrightconfig.json create mode 100644 setup.py diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..2616529 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,16 @@ +.git +.venv +__pycache__ +*.pyc +traces/ +latte-traces/ +debug/replay/ +debug_pub/replay/ +*.mp4 +*.jpg +*.png +pytest_temp/ +/scratch/ +android_world* +/dist/ +node_modules/ diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..4544544 --- /dev/null +++ b/.env.example @@ -0,0 +1,19 @@ +# ============================================================================== +# โ˜• Latte Environment Configuration +# ============================================================================== + +# 1. Primary LLM Provider API Key (Required for Agent Planning & Reasoning) +GEMINI_API_KEY=your_gemini_api_key_here + +# 2. Google Cloud Vision OCR API Key (Required for fast text perception & XML fusion) +OCR_API_KEY=your_google_cloud_vision_api_key_here + +# Optional: Alternative LLM Providers +# OPENAI_API_KEY=your_openai_api_key_here +# ANTHROPIC_API_KEY=your_anthropic_api_key_here +# OPEN_ROUTER_API_KEY=your_openrouter_api_key_here +# XAI_API_KEY=your_xai_api_key_here + +# Optional: Target Android Device Serial (Leave commented for auto-detection) +# ADB_DEVICE_SERIAL=emulator-5554 + diff --git a/.gcloudignore b/.gcloudignore new file mode 100644 index 0000000..4100e79 --- /dev/null +++ b/.gcloudignore @@ -0,0 +1,17 @@ +.gcloudignore +.git +.venv/ +__pycache__/ +*.pyc +traces/ +latte-traces/ +debug/replay/ +debug_pub/replay/ +*.mp4 +*.jpg +*.png +pytest_temp/ +/scratch/ +android_world* +/dist/ +!frontend/dist/** diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..b7580f5 --- /dev/null +++ b/.gitignore @@ -0,0 +1,109 @@ +# ============================================================================== +# โ˜• Latte Framework - Master Git Ignore Configuration +# ============================================================================== + +# --- 1. Environment Variables & Virtual Environments (Recursive) --- +.env +.env.* +!.env.example +uv.lock +**/.venv/ +**/venv/ +**/ENV/ +**/env/ +**/.pyenv/ +.python-version + +# --- 2. Node & Web Frontend Dependencies / Artifacts (Recursive) --- +**/node_modules/ +**/node_modules +**/.angular/ +**/.angular/cache/ +npm-debug.log* +yarn-debug.log* +yarn-error.log* + +# --- 3. Python Caches, Bytecode & Build Artifacts (Recursive) --- +**/__pycache__/ +**/*.py[cod] +**/*$py.class +**/*.so +**/*.pyd +**/*.c +**/build/ +**/dist/ +**/*.egg-info/ +**/.eggs/ +develop-eggs/ +downloads/ +eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +.installed.cfg +*.egg +MANIFEST + +# --- 4. Testing, Coverage & Linters (Recursive) --- +**/.pytest_cache/ +**/.ruff_cache/ +**/.mypy_cache/ +**/.pytype/ +.coverage +.coverage.* +.cache +coverage.xml +nosetests.xml +htmlcov/ +pytest_temp/ +tests_tmp/ +tests/**/outputs/ + +# --- 5. Runtime Logs, Spans, Replay & Traces (Recursive) --- +traces/ +traces_test/ +latte-traces/ +data_engine.db* +**/replay/outputs/ +**/replay/data/ +**/*.log +**/*.log.* +**/*.out +**/*.tmp +.latte_ipc_port +.latte_paused +.latte/ +.artemis_ipc_port +.artemis_paused +.artemis/ +*.bak +live_screenshot* +recording.mp4 +recording.mkv +screenshot.jpg +screen.png + +# --- 6. Workspace Scratch, Dumps & Temporary Experimentation --- +scratch/ +artifacts/ +ignores/ +*.docx +*gdoc.html +*.bug.log + +# --- 7. AI Agent Workspace & IDE Metadata --- +.agents/ +_agents/ +jetski/ +.jetski* +.benchmarks/ +.idea/ +.vscode/ +*.swp +*.swo +*~ +.DS_Store +Thumbs.db diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..9264d43 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,103 @@ +# How to Contribute + +We'd love to accept your patches and contributions to ARTEMIS! There are just a few small guidelines you need to follow. + +--- + +## Contributor License Agreement (CLA) + +Contributions to this project must be accompanied by a Contributor License Agreement (CLA). You (or your employer) retain the copyright to your contribution; this simply gives us permission to use and redistribute your contributions as part of the project. + +- Head over to to check your current agreements on file or to sign a new one. +- If you are an individual contributing original source code, you must sign the **Individual CLA**. +- If you are contributing on behalf of a company or organization, a **Corporate CLA** must be executed by your employer. +- You generally only need to submit a CLA once, so if you have already submitted one for another Google project, you do not need to sign it again. + +--- + +## Community Guidelines + +This project adheres to [Google's Open Source Community Guidelines](https://opensource.google/conduct/). By participating, all contributors and maintainers are expected to uphold these principles and maintain a welcoming, inclusive, and professional environment. + +--- + +## Code Reviews + +All submissions, including submissions by project maintainers, require review. We use standard GitHub pull requests for this purpose: + +1. Consult [GitHub Help](https://docs.github.com/articles/about-pull-requests/) for information on using pull requests. +2. Maintainers will review your changes, offer feedback, and approve when ready. +3. Once approved and all CI checks pass, a project maintainer will merge your PR. + +--- + +## Development & Contribution Workflow + +### 1. Setting Up Your Development Environment + +- **Prerequisites**: + - Python **3.12+** + - Android SDK Platform Tools (`adb` available in `PATH`) + - [`uv`](https://docs.astral.sh/uv/) (recommended) or `pip` + +- **Clone & Install**: + ```bash + git clone https://github.com//artemis.git + cd artemis + + # Install all dependencies and setup pre-commit hooks + make setup + ``` + *(Alternatively: `uv sync --dev` and `uv run pre-commit install`)* + +### 2. Making Changes + +- Create a descriptive topic branch off `main`: + ```bash + git checkout -b feat/my-new-feature + # or + git checkout -b fix/issue-description + ``` + +- Follow the project's coding and style conventions: + - **Absolute Imports**: Always use absolute imports (`from artemis.sdk import Agent`). Relative imports are disallowed. + - **Typing**: Use comprehensive type annotations on all public functions, classes, and methods. + - **Testing**: Add or update test cases under `tests/` for any modified or new behavior. + +### 3. Local Verification + +Before creating a pull request, ensure all linters, type checks, and tests pass: + +```bash +# Format code (Ruff) +make format + +# Lint code (Ruff) +make lint + +# Run static type checker (Pyright) +make typecheck + +# Run test suite (Pytest) +make test +``` + +### 4. Submitting a Pull Request + +1. Push your branch to your fork: + ```bash + git push origin feat/my-new-feature + ``` +2. Open a Pull Request against the `main` branch of the ARTEMIS repository. +3. Complete the PR template, describing: + - The motivation behind the change. + - Any related issue numbers (e.g., `Fixes #42`). + - How the change was tested (device, emulator, unit tests). +4. Ensure the CLA check automatically passes on your PR. + +--- + +## Reporting Bugs and Feature Requests + +- **Bug Reports**: Open an issue describing the bug, including steps to reproduce, device/emulator specifications, Android version, logs, and screenshots if applicable. +- **Feature Requests**: Open an issue detailing the use case, proposed API or workflow, and alternatives considered. diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..03c8762 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,67 @@ +FROM python:3.12-slim + +WORKDIR /app + +# Install necessary build tools +RUN apt-get update && apt-get install -y --no-install-recommends \ + build-essential \ + curl \ + git \ + adb \ + && rm -rf /var/lib/apt/lists/* + +# Copy application source code and apps +COPY artemis/ /app/artemis/ +COPY apps/ /app/apps/ +COPY config/ /app/config/ +COPY pyproject.toml setup.py README.md LICENSE /app/ + +# Install server and core dependencies +RUN pip install --no-cache-dir \ + uvicorn \ + fastapi \ + pydantic \ + pydantic-settings \ + python-dotenv \ + sseclient-py \ + websockets \ + adbutils \ + jinja2 \ + colorama \ + psutil \ + httpx \ + retry \ + google-auth \ + langchain \ + langchain-core \ + langchain-community \ + langchain-google-genai \ + langchain-google-vertexai \ + langchain-mcp-adapters \ + google-genai \ + langgraph \ + typer \ + pillow \ + packaging \ + uiautomator2 \ + opencv-python-headless \ + uuid-utils \ + posthog \ + mcp \ + inquirer \ + rich \ + requests \ + imageio-ffmpeg \ + langchain-openai + +# Install artemis package locally (no-deps to prevent pulling unnecessary heavy optional libraries) +RUN pip install --no-cache-dir --no-deps -e /app + +ENV PYTHONUNBUFFERED=1 +ENV PYTHONPATH=/app +ENV ARTEMIS_APP_DIR=/app +ENV PORT=8080 + +EXPOSE 8080 + +CMD ["uvicorn", "apps.admin_console.server:app", "--host", "0.0.0.0", "--port", "8080"] diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..ba26287 --- /dev/null +++ b/LICENSE @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + https://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + 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 + + https://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. diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..2a1544b --- /dev/null +++ b/Makefile @@ -0,0 +1,56 @@ +.PHONY: help test install setup clean precommit-install precommit lint format typecheck + +help: ## Show this help message + @echo 'Usage: make [target]' + @echo '' + @echo 'Available targets:' + @awk 'BEGIN {FS = ":.*?## "} /^[a-zA-Z_-]+:.*?## / {printf " %-20s %s\n", $$1, $$2}' $(MAKEFILE_LIST) + +test: ## Run all tests + @echo "๐Ÿงช Running all tests..." + @uv run pytest -v + +install: ## Install dependencies + @echo "๐Ÿ“ฆ Installing dependencies..." + @uv sync --dev + +setup: ## Setup project (install dependencies + pre-commit hooks) + @echo "๐Ÿš€ Setting up project..." + @$(MAKE) install + @$(MAKE) precommit-install + @echo "" + @echo "โœ… Setup complete! You're ready to start developing." + +lint: ## Run linting checks + @echo "๐Ÿ” Running linting checks..." + @uv run ruff format --check + @uv run ruff check + +format: ## Format code + @echo "โœจ Formatting code..." + @uv run ruff format + @uv run ruff check --fix + +typecheck: ## Run type checking + @echo "๐Ÿ” Running type checks..." + @uv run pyright + +precommit-install: ## Install pre-commit hooks + @echo "๐Ÿ”ง Installing pre-commit hooks..." + @uv run pre-commit install + @echo "โœ… Pre-commit hooks installed successfully!" + @echo "" + @echo "Pre-commit will now run automatically on every commit." + @echo "To run manually: make precommit" + +precommit: ## Run pre-commit hooks manually on all files + @echo "๐Ÿ” Running pre-commit checks..." + @uv run pre-commit run --all-files + +clean: ## Clean up generated files and caches + @echo "๐Ÿงน Cleaning up..." + @rm -rf .pytest_cache + @rm -rf .ruff_cache + @rm -rf **/__pycache__ + @find . -type d -name "*.egg-info" -exec rm -rf {} + 2>/dev/null || true + @echo "โœ… Cleanup complete" diff --git a/README.md b/README.md new file mode 100644 index 0000000..8113f16 --- /dev/null +++ b/README.md @@ -0,0 +1,399 @@ +# ARTEMIS โ˜•๐Ÿ“ฑ + +**Autonomous Multimodal Android Agent & Mobile UI Automation Platform** + +[![Python 3.12+](https://img.shields.io/badge/python-3.12+-blue.svg)](https://www.python.org/downloads/) +[![License: Apache-2.0](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](LICENSE) +[![Framework: LangGraph](https://img.shields.io/badge/Orchestration-LangGraph-orange.svg)](https://github.com/langchain-ai/langgraph) +[![Model: Google Gemini](https://img.shields.io/badge/Powered%20by-Google%20Gemini-4285F4.svg)](https://ai.google.dev/) + +ARTEMIS is a next-generation, production-grade autonomous agent and testing framework engineered for Android devices and emulators. Combining state-of-the-art multimodal Large Vision-Language Models (VLMs) with hierarchical closed-loop planning, resilient self-healing, and deep Android subsystem integration (ADB, UIAutomator2, accessibility trees), ARTEMIS enables autonomous UI navigation, robust automated testing, and intelligent device diagnostics. + +--- + +## ๐Ÿ“‘ Table of Contents + +- [Key Features](#-key-features) +- [Architecture Overview](#-architecture-overview) +- [Prerequisites](#-prerequisites) +- [Installation & Setup](#-installation--setup) +- [Quick Start](#-quick-start) + - [Command Line Interface (CLI)](#command-line-interface-cli) + - [Python SDK](#python-sdk) +- [Execution Profiles: Flash vs. Pro](#-execution-profiles-flash-vs-pro) +- [Model Context Protocol (MCP) Integration](#-model-context-protocol-mcp-integration) +- [Configuration Guide](#-configuration-guide) + - [Using Non-Google LLMs](#using-non-google-llms) +- [Trace Inspection & Observability](#-trace-inspection--observability) +- [Development & Testing](#-development--testing) +- [License](#-license) + +--- + +## โœจ Key Features + +- **Dual-Engine Execution**: + - **ARTEMIS Flash**: High-speed, reactive single-agent loop for deterministic, low-latency tasks (< 35 steps). + - **ARTEMIS Pro**: Multi-agent closed-loop graph orchestration (Planner $\rightarrow$ Operator $\rightarrow$ Failure Analyzer $\rightarrow$ Summarizer) powered by LangGraph. +- **Multimodal Perception**: + - Combined visual screenshot grounding, real-time OCR, accessibility hierarchy (XML) parsing, and dynamic video analysis. +- **Resilient Locating Strategy ("Dynamic-First, Coordinate-Fallback")**: + - Prioritizes semantic IDs, XPath, and accessibility text matching to withstand UI layout drift; seamlessly falls back to calibrated absolute/relative coordinates. +- **Autonomous Self-Healing**: + - Built-in failure analyzer and safety net detect blocked flows, system dialogs, app crashes, or unexpected states and auto-recover in real time. +- **App Locking & Session Boundaries**: + - Restricts execution context to target application packages with automatic foreground recovery. +- **Model Context Protocol (MCP) Native**: + - Exposes Android automation, state inspection, and trace querying as MCP tools for IDEs and AI orchestrators. +- **Structured Pydantic Outputs**: + - Automatically parses agent observations and findings into validated Pydantic schemas. + +--- + +## ๐Ÿ› Architecture Overview + +ARTEMIS operates on a closed-loop multi-agent architecture: + +``` + โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” + โ”‚ User Goal / Task โ”‚ + โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ + โ”‚ + โ–ผ + โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” + โ”‚ Planner Agent โ”‚ + โ”‚ (Hierarchical Task Breakdown) โ”‚ + โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ + โ”‚ + โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ดโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” + โ–ผ โ–ผ + โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” + โ”‚ Operator Agent โ”‚ โ”‚ Diagnostic Agents โ”‚ + โ”‚ (Screen VLM + XML Grounding) โ”‚ โ”‚ (Log Reader, Video Analyzer, โ”‚ + โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ โ”‚ Object Detector, Explorer) โ”‚ + โ”‚ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ + โ–ผ + โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” + โ”‚ Execution & Device Control โ”‚ + โ”‚ (ADB / UIAutomator2 / Taps) โ”‚ + โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ + โ”‚ + [Action Success / State Change?] + โ”‚ โ”‚ + (Yes) (No / Blocked) + โ”‚ โ”‚ + โ–ผ โ–ผ + โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” + โ”‚ Next Step / Complete โ”‚ โ”‚ Failure Analyzer & Safety Net โ”‚ + โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ โ”‚ (Popup dismissal, Retry, Self-healing) โ”‚ + โ”‚ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ + โ–ผ โ”‚ + โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”‚ + โ”‚ Summarizer / Output โ”‚โ—„โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ + โ”‚ (Reports, Traces) โ”‚ + โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ +``` + +--- + +## ๐Ÿ”ง Prerequisites + +1. **Python**: Version `3.12` or higher. +2. **Android SDK / ADB**: + - `adb` must be installed and accessible in your system `PATH`. + - Verify with `adb devices`. +3. **Android Device or Emulator**: + - A physically connected Android device with **USB Debugging** enabled, or + - An active Android Emulator (AVD) running on your local machine. +4. **Screen Recording & Video Tools** *(optional for `--with-video-recording-tools`)*: + - **`scrcpy`**: Required for live Android device screen video recording (`apt install scrcpy` / `brew install scrcpy`). + - **`ffmpeg`**: Required for video transcoding, trimming, and audio extraction (`apt install ffmpeg` / `brew install ffmpeg`). + +--- + +## ๐Ÿš€ Installation & Setup + +### 1. Clone the Repository + +```bash +git clone https://github.com/google/artemis.git +cd artemis +``` + +### 2. Install Dependencies + +Using [`uv`](https://docs.astral.sh/uv/) (recommended): + +```bash +# Install dependencies into a virtual environment +uv sync --dev +``` + +Or using standard `pip`: + +```bash +python3 -m venv .venv +source .venv/bin/activate +pip install -e ".[dev]" +``` + +### 3. Quick Setup (10 Seconds) + +Run the interactive setup wizard to configure your API key and connect your device: + +```bash +artemis init +``` + +Verify your environment and connected devices anytime with the doctor diagnostic tool: + +```bash +artemis doctor +``` + +*(Alternatively, copy `.env.example` to `.env` and set `GEMINI_API_KEY` for LLM reasoning and `OCR_API_KEY` for Vision perception.)* + +--- + +## ๐Ÿ’ก Quick Start + +### Command Line Interface (CLI) + +ARTEMIS installs a unified command-line tool `artemis`: + +```bash +# 1. Environment & Device Diagnostics +artemis doctor + +# 2. Interactive Setup Wizard +artemis init + +# 3. Basic task execution (ARTEMIS Pro by default) +artemis run "Open YouTube and search for Lo-Fi Hip Hop" + +# 4. Fast reactive profile (ARTEMIS Flash) +artemis run "Open Settings and verify Airplane Mode is turned off" --profile flash + +# 5. Execute with trace recording and structured output +artemis run "Open Maps, search for coffee shops, and extract the top 3 results" \ + --test-name "coffee_search_test" \ + --traces-path "./traces" \ + --output-description '{"places": "list of shop names", "ratings": "list of rating scores"}' + +# 6. Enable dynamic video recording analysis +artemis run "Play the latest video in the feed and verify playback" --with-video-recording-tools +``` + +### Python SDK + +Use the fluent SDK to integrate ARTEMIS into your test suites or automation workflows: + +```python +import asyncio +from pydantic import BaseModel, Field +from artemis.sdk import Agent + + +class VideoSearchResult(BaseModel): + query: str = Field(..., description="The search term used") + top_video_title: str = Field(..., description="Title of the first video result") + channel_name: str = Field(..., description="Channel name of the first video") + + +async def main(): + # 1. Initialize the agent + agent = Agent() + await agent.init() + + try: + # 2. Build task specification + task = ( + agent.new_task("Search for 'Pixel 9 Pro' and get the first video details") + .with_name("youtube_search_verification") + .with_locked_app_package("com.google.android.youtube") + .with_output_format(VideoSearchResult) + .with_max_steps(50) + .build() + ) + + # 3. Execute task + result: VideoSearchResult = await agent.run_task(request=task) + + if result: + print("Search Results:") + print(f"- Query: {result.query}") + print(f"- First Video: {result.top_video_title}") + print(f"- Channel: {result.channel_name}") + else: + print("Task did not complete successfully.") + + finally: + # 4. Clean up agent session & device connections + await agent.clean() + + +if __name__ == "__main__": + asyncio.run(main()) +``` + +--- + +## โšก Execution Profiles: Flash vs. Pro + +| Feature | ARTEMIS Flash (`--profile flash`) | ARTEMIS Pro (`--profile default`) | +| :--- | :--- | :--- | +| **Model Loop** | Single reactive agent (Observe-Think-Act) | Multi-node LangGraph orchestration | +| **Best For** | Short, deterministic tasks (< 35 steps) | Deep, complex, multi-stage workflows | +| **Step Latency** | ~3โ€“5 seconds / step | ~15โ€“30 seconds / turn (with planning & validation) | +| **Planning** | Direct goal execution | Hierarchical multi-step plan decomposition | +| **Diagnostics** | Basic error reporting | Deep diagnostics, log analysis & self-healing | + +--- + +## ๐Ÿ”Œ Model Context Protocol (MCP) Integration + +ARTEMIS includes built-in MCP servers and tools to allow AI IDEs (such as Claude Desktop, Cursor, and Gemini Jetski) to control Android devices directly: + +| Tool Name | Description | +| :--- | :--- | +| `mobile_run_task` | Launches an autonomous background automation task on the connected device. | +| `mobile_manage_task` | Inspects status, terminates tasks, or injects real-time steering instructions. | +| `mobile_get_device_state` | Retrieves real-time screenshots or simplified accessibility UI trees. | +| `mobile_inspect_trace` | Inspects step-by-step screenshots, visual action overlays, and execution logs. | + +--- + +## โš™๏ธ Configuration Guide + +All model configurations, thinking budgets, agent capabilities, and tool filtering are consolidated into a single self-explanatory configuration file: [`config/artemis.jsonc`](config/artemis.jsonc). + +```jsonc +{ + // 1. Global Default Model (Applied across all sub-agents unless overridden) + "default": { + "provider": "google", + "model": "gemini-3.6-flash", + "thinking_level": "medium", + "fallback": { + "provider": "google", + "model": "gemini-3.5-flash" + } + }, + + // 2. Built-in Presets + "presets": { + "gemini-flagship": { "provider": "google", "model": "gemini-3.6-flash" }, + "openai-gpt4o": { "provider": "openai", "model": "gpt-4o" }, + "cost-saving": { "provider": "google", "model": "gemini-3.5-flash-lite" }, + "local-ollama": { "provider": "openai", "model": "llama3.2-vision" } + }, + + // 3. Agent & Sub-Agent Node Customizations (Optional overrides) + "nodes": { + // ๐Ÿง  Planner: High-level goal decomposition and strategic step planning + "planner": { "thinking_level": "medium" }, + + // ๐ŸŽฏ Operator: Screen perception, element interaction, and direct execution + "operator": { "thinking_level": "medium", "include_thoughts": true }, + + // ๐Ÿ›ก๏ธ Validator & Failure Analyzer: Self-healing and auto-recovery + "validator_failure_analyzer": { "thinking_level": "medium" } + }, + + // 4. Agent Capabilities & Tool Filtering + "agent": { + "explorer_versions": { "operator": "flash", "validator": "flash" }, + "blacklisted_tools": { + "explorer": ["ask_image_processor", "get_ocr_list", "inspect_region"] + } + } +} +``` + +> [!TIP] +> **Zero-Friction Global Change**: To switch the entire platform from Gemini to GPT-4o or Claude, simply edit the `"default"` block once in `config/artemis.jsonc`. All 15 sub-agents automatically inherit the new model! + +### Using Non-Google LLMs + +ARTEMIS supports multi-provider LLM backends via standard LangChain integrations (Google Gemini, OpenAI, OpenRouter, xAI / Grok, or Google Vertex AI). + +#### Supported Providers & API Key Setup + +| Provider | Provider Key | Environment Variable | Base URL Override | Example Models | +| :--- | :--- | :--- | :--- | :--- | +| **Google Gemini** | `google` *(default)* | `GEMINI_API_KEY` | N/A | `gemini-3.6-flash`, `gemini-3.5-flash` | +| **Google Vision OCR** | `ocr` | `OCR_API_KEY` | N/A | Google Cloud Vision Text Detection | +| **Google Vertex AI** | `vertexai` | Google ADC (`gcloud auth application-default login`) | N/A | `gemini-3.6-flash` | +| **OpenAI** | `openai` | `OPENAI_API_KEY` | `OPENAI_BASE_URL` | `gpt-4o`, `gpt-4o-mini`, `o3` | +| **OpenRouter** | `openrouter` | `OPEN_ROUTER_API_KEY` | N/A | `anthropic/claude-3.5-sonnet`, `meta-llama/llama-3.3-70b-instruct` | +| **xAI (Grok)** | `xai` | `XAI_API_KEY` | N/A | `grok-2-vision-1212`, `grok-beta` | + +> [!NOTE] +> ARTEMIS uses multimodal visual grounding (screenshots) for UI navigation. When assigning non-Google models to vision-dependent agents (such as `operator`), ensure the selected model supports image/vision inputs. + +#### Switching Provider in `config/artemis.jsonc` + +To use OpenAI or Claude globally, simply change the `default` section in [`config/artemis.jsonc`](config/artemis.jsonc): + +```jsonc +{ + "default": { + "provider": "openai", + "model": "gpt-4o", + "fallback": { + "provider": "openai", + "model": "gpt-4o-mini" + } + } +} +``` + +#### Custom / Local OpenAI-Compatible Endpoints (Ollama, vLLM, LocalAI) + +To connect ARTEMIS to local or self-hosted LLM endpoints (e.g., Ollama or vLLM), set `OPENAI_BASE_URL` in `.env` and specify `openai` as the provider in your configuration: + +```bash +# Route OpenAI provider calls to local Ollama / vLLM endpoint +OPENAI_API_KEY=ollama # dummy key if required by client +OPENAI_BASE_URL=http://localhost:11434/v1 +``` + +--- + +## ๐Ÿ“Š Trace Inspection & Observability + +Every task execution can record rich execution artifacts stored in `./traces//`: +- **`notes/`**: High-level planner notes, step summaries, and generated reports (`output.md`). +- **`screenshots/`**: + - `before_screenshot.jpg`: Initial screen observed by the agent. + - `action_overlay_screenshot.jpg`: Screen with visual tap markers (red circles) and swipe vectors (red arrows). +- **`logs/`**: Raw framework, operator, and failure analyzer diagnostic logs. + +--- + +## ๐Ÿ›  Development & Testing + +ARTEMIS provides standard `make` targets for development: + +```bash +# Setup environment and pre-commit hooks +make setup + +# Run test suite +make test + +# Run code formatters (ruff) +make format + +# Run linter checks +make lint + +# Run static type checker (pyright) +make typecheck +``` + +--- + +## ๐Ÿ“„ License + +This project is licensed under the [Apache License 2.0](LICENSE). diff --git a/apps/README.md b/apps/README.md new file mode 100644 index 0000000..486c62f --- /dev/null +++ b/apps/README.md @@ -0,0 +1,26 @@ +# ๐Ÿ“ฑ Artemis Applications Ecosystem (`apps/`) + +This directory houses the two primary applications and user interfaces of the Artemis ecosystem. + +--- + +## ๐Ÿ“‚ Subsystems Overview + +``` +apps/ +โ”œโ”€โ”€ ๐Ÿ› ๏ธ admin_console/ # [Interface 1] Full-fidelity Trace Inspector & Task Control Console (FastAPI + Web) +โ””โ”€โ”€ โœจ showcase_ui/ # [Interface 2] Aesthetic User Presentation & Interaction Workspace (Angular 19) +``` + +--- + +## ๐Ÿ” Detailed Comparison of the Two UIs + +| Dimension | ๐Ÿ› ๏ธ Admin & Trace Console (`admin_console`) | โœจ Showcase & Workspace UI (`showcase_ui`) | +| :--- | :--- | :--- | +| **Primary Audience** | Developers, Test Engineers, System Admins | End Users, Customers, Product Demo Audiences | +| **UI Paradigm** | High-density, full-fidelity tree logs, step-by-step inspector | Dynamic Aurora Glow, Glassmorphism, dual-pane stream | +| **Core Capabilities** | โ€ข Full hierarchical Call Tree & LLM thinking inspection
โ€ข Pre/post action screenshots & visual overlay verification
โ€ข Video playback & step replay engine (`/api/replay`)
โ€ข Queue tasks & live state toggle | โ€ข Real-time Agent thought & action progress streaming
โ€ข Interactive natural language chat sidebar
โ€ข Aesthetic visual representation of mobile autonomy | +| **Tech Stack** | Python (FastAPI) + Vanilla Responsive Web | TypeScript, Angular 19, SCSS | +| **Start Command** | `uv run uvicorn apps.admin_console.server:app --port 8000 --reload` | `cd apps/showcase_ui && npm start` (Port 4200) | + diff --git a/apps/__init__.py b/apps/__init__.py new file mode 100644 index 0000000..3510f47 --- /dev/null +++ b/apps/__init__.py @@ -0,0 +1 @@ +"""Artemis / Latte Applications Package.""" diff --git a/artemis/__init__.py b/artemis/__init__.py new file mode 100644 index 0000000..a2f5b95 --- /dev/null +++ b/artemis/__init__.py @@ -0,0 +1,41 @@ +# 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. + +"""โ˜• ARTEMIS: Autonomous Multimodal Android Agent & Testing Framework.""" + +__version__ = "0.2.0" + +from artemis.interfaces.sdk.client import ArtemisClient +from artemis.interfaces.sdk.task import StreamEvent, StreamEventType, Task, TaskResult +from artemis.sdk.agent import Agent +from artemis.sdk.builders import Builders +from artemis.agents.base import BaseAgent, AgentRegistry +from artemis.drivers.base import BaseDeviceDriver +from artemis.tools.base import artemis_tool, ToolRegistry + +__all__ = [ + "__version__", + "ArtemisClient", + "StreamEvent", + "StreamEventType", + "Agent", + "Builders", + "Task", + "TaskResult", + "BaseAgent", + "AgentRegistry", + "BaseDeviceDriver", + "artemis_tool", + "ToolRegistry", +] diff --git a/artemis/__main__.py b/artemis/__main__.py new file mode 100644 index 0000000..3199a37 --- /dev/null +++ b/artemis/__main__.py @@ -0,0 +1,18 @@ +# 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. + +from artemis.interfaces.cli.main import app + +if __name__ == "__main__": + app() diff --git a/artemis/agents/__init__.py b/artemis/agents/__init__.py new file mode 100644 index 0000000..bd79617 --- /dev/null +++ b/artemis/agents/__init__.py @@ -0,0 +1,37 @@ +# 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. + +"""Modular Agent Framework for ARTEMIS.""" + +from artemis.agents.base import ( + AgentConfig, + AgentRegistry, + AgentResponse, + BaseAgent, +) +from artemis.agents.operator.operator_agent import OperatorAgent +from artemis.agents.planner.planner_agent import PlannerAgent +from artemis.agents.summarizer.summarizer_agent import SummarizerAgent +from artemis.agents.validator.failure_analyzer_agent import FailureAnalyzerAgent + +__all__ = [ + "BaseAgent", + "AgentConfig", + "AgentResponse", + "AgentRegistry", + "PlannerAgent", + "OperatorAgent", + "SummarizerAgent", + "FailureAnalyzerAgent", +] diff --git a/artemis/agents/base.py b/artemis/agents/base.py new file mode 100644 index 0000000..67ddc8f --- /dev/null +++ b/artemis/agents/base.py @@ -0,0 +1,146 @@ +# 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. + +"""Base Agent Protocol & Lifecycle Framework for ARTEMIS. + +Provides a unified interface and lifecycle management for all specialized +multimodal agents (Planner, Operator, FailureAnalyzer, Explorer, Summarizer, etc.). +""" + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any + +from pydantic import BaseModel, Field + +from artemis.drivers.factory import get_driver +from artemis.utils.logger import get_logger + +if TYPE_CHECKING: + pass + +logger = get_logger(__name__) + + +class AgentConfig(BaseModel): + """Configuration options for an individual agent.""" + + model_name: str | None = Field(default=None, description="LLM/VLM model override") + temperature: float = Field(default=0.0, description="Sampling temperature") + max_tokens: int = Field(default=4096, description="Maximum output tokens") + system_prompt_override: str | None = Field( + default=None, description="Custom system prompt override" + ) + enabled_tools: list[str] = Field( + default_factory=list, description="Tool names enabled for this agent" + ) + + +class AgentResponse(BaseModel): + """Standardized structured response from an agent invocation.""" + + status: str = Field( + default="success", description="Execution status ('success', 'retry', 'fail', 'done')" + ) + thought: str | None = Field(default=None, description="Reasoning and internal monologue") + action_taken: str | None = Field(default=None, description="Action or tool called") + output: Any = Field(default=None, description="Output payload or state updates") + metadata: dict[str, Any] = Field( + default_factory=dict, description="Diagnostic and profiling metadata" + ) + + +class BaseAgent(ABC): + """Abstract Base Class defining the standard lifecycle for all ARTEMIS agents.""" + + def __init__( + self, + name: str, + role: str, + description: str, + config: AgentConfig | None = None, + ctx: Any | None = None, + ): + self.name = name + self.role = role + self.description = description + self.config = config or AgentConfig() + self.ctx: Any | None = ctx + self._is_initialized = False + + @property + def driver(self) -> Any | None: + """Helper to get the active device driver from the context.""" + if self.ctx: + return get_driver(self.ctx) + return None + + async def initialize(self, ctx: Any) -> None: + """Lifecycle hook: Called once when the agent is loaded into the workflow.""" + self.ctx = ctx + self._is_initialized = True + logger.info(f"Agent '{self.name}' ({self.role}) initialized.") + + @abstractmethod + async def process(self, state: Any) -> dict[str, Any]: + """Main execution hook: Receives state, reasons, interacts with tools, and returns state updates.""" + ... + + async def cleanup(self) -> None: + """Lifecycle hook: Called upon session or workflow termination.""" + self._is_initialized = False + logger.debug(f"Agent '{self.name}' cleanup complete.") + + def __repr__(self) -> str: + return f"<{self.__class__.__name__} name={self.name!r} role={self.role!r}>" + + +class AgentRegistry: + """Central registry and factory for ARTEMIS agents.""" + + _registry: dict[str, type[BaseAgent]] = {} + + @classmethod + def register(cls, name: str | None = None): + """Decorator to register an agent class with a unique identifier.""" + + def decorator(subclass: type[BaseAgent]) -> type[BaseAgent]: + reg_name = name or subclass.__name__.lower() + cls._registry[reg_name] = subclass + return subclass + + return decorator + + @classmethod + def get_agent_class(cls, name: str) -> type[BaseAgent] | None: + return cls._registry.get(name.lower()) + + @classmethod + def create_agent( + cls, + name: str, + ctx: Any | None = None, + config: AgentConfig | None = None, + **kwargs: Any, + ) -> BaseAgent: + """Instantiates an agent by registered name.""" + agent_cls = cls.get_agent_class(name) + if not agent_cls: + raise KeyError( + f"Agent '{name}' is not registered. Available: {list(cls._registry.keys())}" + ) + return agent_cls(ctx=ctx, config=config, **kwargs) + + @classmethod + def list_agents(cls) -> list[str]: + return list(cls._registry.keys()) diff --git a/artemis/agents/checker/checker.json b/artemis/agents/checker/checker.json new file mode 100644 index 0000000..cba8b05 --- /dev/null +++ b/artemis/agents/checker/checker.json @@ -0,0 +1,4 @@ +{ + "checker_system": "You are a precise and reliable verification expert (Checker). Your task is to determine whether the Operator has successfully completed the active subgoal by combining visual evidence from the current screenshot with data stored in persistent notes.", + "checker_prompt": "# Inputs\n- Active Subgoal: {{ subgoal }}\n- Execution History:\n{{ plan_and_history }}\n- Dialogue History: {{ dialogue_history }}\n- Visual State: The provided screenshot.\n\n# Guidelines\n- **Visual + Data Verification**: You must combine the screenshot with note contents. If the subgoal involves extracting, saving, or recording information (e.g., \"save contacts\", \"record address\"), you **MUST** verify that the content actually exists in the note. Do not assume success just because the information is visible on the screen. If the signs of success are not obvious, you can check the history, especially the most recent steps, to see if the operator performed their own check and whether that check was reasonable, to determine if the subgoal was truly successful.\n- **Visual Priority (Anti-History-Bias)**: Your primary ground truth is the visual evidence on the **Current Screenshot**. The `Execution History` lists previous successes or failures, which may be outdated due to dynamic background recovery or self-healing. If the screenshot demonstrably shows the active subgoal's target state is achieved, you **MUST** mark `success: true`, regardless of previous failed steps. Trust what you see, not past FAILED texts.\n- **Constraint Verification**: For subgoals containing specific parameters (such as names, brands, topics, events, numbers, or dates), confirm that the visible text precisely reflects these specified details. A successful state is defined by a clear, direct match with the subgoal's constraints.\n- **Realistic Expectations**: You cannot scroll the screen. If a subgoal requires saving all items but only some are visible on the screen, verify that at least the **visible items** are recorded in the note. As long as the visible part is recorded and there are no obvious errors, consider it a success. Do not fail the task simply because off-screen content cannot be verified.\n- **Success Criteria**: Set `success` to `true` if the core objective of the subgoal is met and it does not block subsequent steps. Set `success` to `false` if the state is incorrect, required data was not saved, or the operator is stuck/looping.\n\n# Output Protocol\nInvoke the `submit_verification` tool with `success` (boolean) and `reason` (string) to submit your final decision. If you need to inspect note data first, use `read_note`." +} diff --git a/artemis/agents/checker/checker.py b/artemis/agents/checker/checker.py new file mode 100644 index 0000000..caaa69a --- /dev/null +++ b/artemis/agents/checker/checker.py @@ -0,0 +1,460 @@ +# 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 asyncio +import hashlib +import json +from pathlib import Path +import re +from typing import Any + +from jinja2 import Template +from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage +from langchain_core.tools import tool + +from pydantic import BaseModel, Field + +from artemis.constants import CHECKER_MAX_ITERATIONS +from artemis.context import ArtemisContext +from artemis.controllers.unified_controller import UnifiedMobileController +from artemis.data_engine.trace import ( + TraceSpan, + trace, + trace_langchain_tool, +) +from artemis.services.llm import get_llm, invoke_llm_with_timeout_message +from artemis.tools.index import get_tool_by_name +from artemis.tools.scratchpad import get_read_note_tool_pure +from artemis.tools.tool_wrapper import ( + get_tool_result_content, + invoke_tool_with_injection, +) +from artemis.utils.logger import get_logger +from artemis.utils.notes import get_note_file_path, get_notes_dir +from artemis.utils.ocr_api import perform_ocr +from artemis.utils.ocr_xml_fusion import ( + _crop_image_remove_status_bar, + _detect_status_bar_height, + _map_coordinates_back, + fuse_ocr_with_xml, +) +from artemis.utils.task_tree import ( + build_plan_and_history, + get_recent_subgoal_hashes, +) +from artemis.utils.verification import ( + append_verification_chat, + get_verification_chat_path, + get_verification_chat_rounds, + read_verification_chat, +) +from artemis.utils.text import safe_extract_text +from artemis.utils.visualization import format_minimal_list_with_elements + +logger = get_logger(__name__) + + +class CheckerResult(BaseModel): + """Structured output contract for Checker verification.""" + + success: bool = Field( + description="Whether the active subgoal was successfully completed based on visual and note evidence." + ) + reason: str = Field( + description="Detailed explanation of the result. If failed, specify what is missing or wrong and give actionable hints for the Operator." + ) + + +@trace(type="agent", name="checker") +async def run_async_check( + ctx: ArtemisContext, + subgoal_text: str, + subgoal_hash: str | None = None, + raw_perception_data: dict[str, Any] | None = None, + latest_ui_hierarchy: list[Any] | None = None, +): + """Runs the checker verification as a background task.""" + logger.info(f"Starting async Checker for: {subgoal_text}") + + if not ctx.data_engine: + logger.warning("DataEngine not available, skipping checker.") + return {"status": "success", "reason": "DataEngine unavailable."} + + notes_dir = get_notes_dir(ctx.data_engine.base_dir) + status_path = notes_dir / "checker_status.json" + + if ctx.execution_setup and ctx.execution_setup.disable_checker: + logger.info("Checker is disabled via flag. Skipping.") + try: + status_path.write_text( + json.dumps({"status": "success", "reason": "Checker disabled."}), + encoding="utf-8", + ) + except Exception: + pass + return {"status": "success", "reason": "Checker disabled."} + + # 1. Set status to running + try: + status_path.write_text( + json.dumps({"status": "running", "subgoal": subgoal_text}), + encoding="utf-8", + ) + logger.info(f"Wrote checker status 'running' to {status_path}") + except Exception as e: + logger.error(f"Failed to write checker status: {e}") + return {"status": "success", "reason": f"Status write error: {e}"} + + # 2. Run Checker logic + try: + # Retrieve perception data: directly reuse upstream perception if available + if isinstance(raw_perception_data, dict) and raw_perception_data.get("screenshot_b64"): + screenshot_b64 = raw_perception_data["screenshot_b64"] + width = raw_perception_data.get("width", 1080) + height = raw_perception_data.get("height", 2400) + fused_xml = latest_ui_hierarchy or raw_perception_data.get("xml_hierarchy") or [] + elif ( + hasattr(ctx, "latest_perception_data") + and isinstance(ctx.latest_perception_data, dict) + and ctx.latest_perception_data + ): + screenshot_b64 = ctx.latest_perception_data.get("screenshot_b64") + width = ctx.latest_perception_data.get("width", 1080) + height = ctx.latest_perception_data.get("height", 2400) + fused_xml = ctx.latest_perception_data.get("latest_ui_hierarchy") or [] + else: + # Fallback to direct controller query + controller = UnifiedMobileController(ctx) + device_data = await controller.get_screen_data() + screenshot_b64 = device_data.base64 + xml_hierarchy = device_data.elements + width = device_data.width + height = device_data.height + + ocr_results = [] + try: + status_bar_height = _detect_status_bar_height(xml_hierarchy, height) + if status_bar_height > 0: + cropped_b64, _, _ = _crop_image_remove_status_bar( + screenshot_b64, status_bar_height + ) + raw_ocr_results = await perform_ocr(cropped_b64) + ocr_results = _map_coordinates_back(raw_ocr_results, status_bar_height) + else: + ocr_results = await perform_ocr(screenshot_b64) + except Exception as e: + logger.error(f"Checker OCR failed, proceeding without OCR: {e}") + + fused_xml = fuse_ocr_with_xml(xml_hierarchy, ocr_results) + + minimal_list, elements, labels = format_minimal_list_with_elements(fused_xml, width, height) + + # Build task tree + task_plan = "No task plan yet." + task_plan_path = get_note_file_path(ctx.data_engine.base_dir, "task_plan") + if task_plan_path.exists(): + try: + task_plan = task_plan_path.read_text(encoding="utf-8") + except Exception as e: + logger.error(f"Failed to read task plan for Checker: {e}") + + steps = ctx.data_engine.get_agent_friendly_steps() if ctx.data_engine else [] + + # Calculate exact subgoal hash from subgoal_text directly to avoid 'default' placeholder + if not subgoal_hash or subgoal_hash == "default": + clean_subgoal = safe_extract_text(subgoal_text).strip() + subgoal_hash = hashlib.md5(clean_subgoal.encode("utf-8")).hexdigest() + + keep_hashes = get_recent_subgoal_hashes(steps, subgoal_hash, ctx.data_engine.base_dir) + + plan_and_history = build_plan_and_history( + task_plan, steps, subgoal_hash, keep_subgoal_hashes=keep_hashes + ) + + chat_path = get_verification_chat_path(ctx.data_engine.base_dir, subgoal_hash) + turns = read_verification_chat(chat_path) + + # Calculate round number + max_op, max_chk = get_verification_chat_rounds(turns) + round_num = max(max_op, max_chk + 1) + if round_num == 0: + round_num = 1 + + # Render dialogue history for prompt + dialogue_lines = [] + for t in turns: + role = "Operator" if t["role"] == "operator" else "Checker" + dialogue_lines.append(f"**{role} (Round {t['round']})**:\n{t['content']}") + dialogue_history = ( + "\n\n".join(dialogue_lines) if dialogue_lines else "No previous dialogue." + ) + + prompts = {} + prompts_path = Path(__file__).parent / "checker.json" + if prompts_path.exists(): + try: + prompts = json.loads(prompts_path.read_text(encoding="utf-8")) + except Exception as e: + logger.error(f"Failed to load checker prompts: {e}") + + prompt_template = prompts.get( + "checker_prompt", + "You are the Checker. Your job is to verify if the current subtask" + " was successfully completed.", + ) + + full_prompt = Template(prompt_template).render( + subgoal=subgoal_text, + plan_and_history=plan_and_history, + dialogue_history=dialogue_history if dialogue_history else "No previous dialogue.", + minimal_list=minimal_list, + ) + + llm = get_llm(ctx=ctx, name="checker") + read_note = get_read_note_tool_pure(ctx) + + @tool + def submit_verification( + success: bool, + reason: str, + ) -> str: + """[TERMINAL] Submit the final verification outcome (success or failure with reason). This completes verification.""" + return f"Verification decision: success={success}, reason={reason}" + + tools_to_bind = [ + trace_langchain_tool(read_note, ctx), + submit_verification, + ] + + content = [ + {"type": "text", "text": full_prompt}, + {"type": "text", "text": "--- Current Screenshot ---"}, + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{screenshot_b64}"}, + }, + ] + + messages = [ + SystemMessage( + content=prompts.get( + "checker_system", + "You are the Checker agent. Your objective is to accurately" + " assess whether the active subgoal was successfully" + " completed.", + ) + ), + HumanMessage(content=content), + ] + + max_iterations = ( + getattr( + ctx.execution_setup, + "checker_max_iterations", + CHECKER_MAX_ITERATIONS, + ) + if ctx.execution_setup + else CHECKER_MAX_ITERATIONS + ) + + success = True + reason = "Checker completed." + + for i in range(max_iterations): + is_final_round = i == max_iterations - 1 + + if is_final_round: + messages.append( + HumanMessage( + content=( + "This is your final iteration; all tools are" + " stripped, and you must provide your final structured verification result." + ) + ) + ) + structured_llm = llm.with_structured_output(CheckerResult) + structured_res: CheckerResult = await invoke_llm_with_timeout_message( + structured_llm.ainvoke(messages) + ) + if isinstance(structured_res, CheckerResult): + success = structured_res.success + reason = structured_res.reason + elif isinstance(structured_res, dict): + success = structured_res.get("success", True) + reason = structured_res.get("reason", "Verification completed.") + break + else: + if i > 0: + messages.append( + HumanMessage( + content=( + "You have not completed verification yet" + f" (iteration {i + 1} of {max_iterations})." + ) + ) + ) + active_llm = llm.bind_tools(tools=tools_to_bind) + + async def run_stream(): + full_response = None + async for chunk in active_llm.astream(messages): + if full_response is None: + full_response = chunk + else: + full_response += chunk + return full_response + + response = await invoke_llm_with_timeout_message(run_stream()) + + # Check if submit_verification was invoked + submit_tc = next( + ( + tc + for tc in response.tool_calls + if tc["name"].endswith("submit_verification") + ), + None, + ) + if submit_tc: + tc_args = dict(submit_tc["args"]) + success = bool(tc_args.get("success", True)) + reason = str(tc_args.get("reason", "Verification submitted.")) + break + + if not response.tool_calls: + # Model answered in plain text; use structured_llm directly + structured_llm = llm.with_structured_output(CheckerResult) + structured_res: CheckerResult = await invoke_llm_with_timeout_message( + structured_llm.ainvoke(messages + [response]) + ) + if isinstance(structured_res, CheckerResult): + success = structured_res.success + reason = structured_res.reason + elif isinstance(structured_res, dict): + success = structured_res.get("success", True) + reason = structured_res.get("reason", "Verification completed.") + break + + messages.append(response) + + async def run_tool(tc): + tool_name = tc["name"].split(":")[-1] if ":" in tc["name"] else tc["name"] + logger.info(f"Checker requested tool: {tool_name}") + + result_str = "" + tool_status = "success" + try: + if ":" in tool_name: + tool_to_run = get_tool_by_name(tool_name, traced_tools) + else: + tool_to_run = next( + (t for t in traced_tools if t.name == tool_name), + None, + ) + if tool_to_run: + args = dict(tc["args"]) + with TraceSpan(name=tool_name, ctx=ctx) as span: + span.payload = {"args": args} + try: + result_obj = await invoke_tool_with_injection( + tool=tool_to_run, + args=args, + tool_call_id=tc["id"], + ) + result_str = get_tool_result_content(result_obj) + span.result = result_str + if result_str.startswith("Error"): + tool_status = "error" + span.status = "failed" + span.error = result_str + except Exception as err: + tool_status = "error" + span.status = "failed" + span.error = str(err) + raise err + else: + result_str = f"Error: Tool {tool_name} not supported" + tool_status = "error" + except Exception as err: + logger.error(f"Error running tool {tool_name}: {err}") + result_str = f"Error running tool {tool_name}: {err}" + tool_status = "error" + + return ToolMessage( + tool_call_id=tc["id"], + content=result_str, + status=tool_status, + ) + + tool_outputs = await asyncio.gather(*(run_tool(tc) for tc in response.tool_calls)) + for tm in tool_outputs: + messages.append(tm) + + logger.info( + f"Async Checker result for '{subgoal_text}': Success={success}, Reason={reason}" + ) + + # Update status file + status_path.write_text( + json.dumps({"status": "success" if success else "failed", "reason": reason}), + encoding="utf-8", + ) + + # Clean up dialogue file if successful + if success: + if chat_path.exists(): + logger.info(f"Checker succeeded. Cleaning up dialogue file: {chat_path}") + chat_path.unlink() + + # If failed, revert status in task_plan.md and write feedback + if not success: + task_plan_path = get_note_file_path(ctx.data_engine.base_dir, "task_plan") + if task_plan_path.exists(): + task_plan = task_plan_path.read_text(encoding="utf-8") + lines = task_plan.split("\n") + + for i in range(len(lines) - 1, -1, -1): + if lines[i].startswith("- ["): + task_content = safe_extract_text(lines[i][5:]).strip() + safe_subgoal = safe_extract_text(subgoal_text).strip() + if task_content == safe_subgoal: + lines[i] = re.sub(r"^(- \[)[x/ ](\])", r"\g<1>/\g<2>", lines[i]) + logger.info( + "Matched top-level subgoal by content." + f" Reverted status to [/]: {subgoal_text}" + ) + break + task_plan_path.write_text("\n".join(lines), encoding="utf-8") + + append_verification_chat( + chat_path, + "checker", + f"Status: Failed.\nReason: {reason}", + round_num, + ) + + return {"status": "success" if success else "failed", "reason": reason} + + except Exception as e: + logger.error(f"Async Checker encountered exception (fail-open): {e}") + status_path.write_text( + json.dumps( + { + "status": "success", + "reason": f"Checker Fail-Open: {e}", + } + ), + encoding="utf-8", + ) + return {"status": "success", "reason": f"Checker Fail-Open: {e}"} diff --git a/artemis/agents/diagnoser/diagnoser.json b/artemis/agents/diagnoser/diagnoser.json new file mode 100644 index 0000000..2df3899 --- /dev/null +++ b/artemis/agents/diagnoser/diagnoser.json @@ -0,0 +1,22 @@ +{ + "system_prompt": [ + "# Diagnoser", + "", + "## Objective", + "You are an expert AI diagnoser. Your objective is to pinpoint the root cause of mobile execution failures by cross-referencing visual evidence (videos/screenshots) with system telemetry (logs). Precisely answer the `Diagnostic Query` based on the provided context.", + "", + "## Core Guidelines", + "- **Evidence-Based**: Base all conclusions strictly on facts from tools and context. Do not assume or hallucinate.", + "- **Laser Focus**: Focus exclusively on the `Diagnostic Query`. Do not investigate unrelated issues.", + "- **No Direct Operations**: Do not perform user actions or modify device state (e.g., tap, swipe, text input, BACK/HOME keys, or force-starting activities). Focus solely on analysis.", + "- **Tool Constraints**: Video analysis is expensive. When calling the video analyzer, you MUST provide a specific, narrow timeframe to minimize cost.", + "- **Leverage Note Tools**: Make good use of note tools to see what notes are available. When diagnosing, open and cross-reference relevant notes to ensure you have grasped all necessary information and context.", + "- **Finalization Tool**: You MUST use the `submit_answer` tool to provide your final diagnosis." + ], + "checker_feedback_template": [ + "--- Checker Feedback ---", + "The Checker agent has provided feedback regarding the verification of the active subgoal:", + "{checker_feedback}", + "" + ] +} \ No newline at end of file diff --git a/artemis/agents/diagnoser/diagnoser.py b/artemis/agents/diagnoser/diagnoser.py new file mode 100644 index 0000000..2fd5239 --- /dev/null +++ b/artemis/agents/diagnoser/diagnoser.py @@ -0,0 +1,575 @@ +# 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 asyncio +import base64 +import json +import os +from pathlib import Path +from typing import Any +import uuid + +from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage + +from artemis.context import ArtemisContext +from artemis.data_engine.trace import ( + CURRENT_TRACE_ID, + TraceSpan, + trace, + trace_langchain_tool, +) +from artemis.graph.state import State +from artemis.services.llm import get_llm, invoke_llm_with_timeout_message +from artemis.tools.command_tool import get_run_short_adb_command_tool +from artemis.tools.diagnoser_submit_answer_tool import get_submit_answer_tool +from artemis.tools.index import get_tool_by_name +from artemis.tools.log_tool import get_analyze_logs_tool +from artemis.tools.mobile.read_hierarchy import get_ui_hierarchy_tool +from artemis.tools.scratchpad import get_list_notes_tool, get_read_note_tool +from artemis.tools.tool_wrapper import ( + get_tool_result_content, + invoke_tool_with_injection, +) +from artemis.tools.video_tool import get_video_analyzer_tool +from artemis.tools.wait_tool import get_wait_tool +from artemis.utils.logger import get_logger +from artemis.utils.task_tree import ( + build_plan_and_history, + get_active_subgoal_hashes, + get_recent_subgoal_hashes, +) + +logger = get_logger(__name__) + + +class Diagnoser: + def __init__(self, ctx: ArtemisContext): + self.ctx = ctx + self.is_device_online = self._is_device_online() + + def _is_device_online(self) -> bool: + try: + if self.ctx.adb_client is None: + return False + if getattr(self.ctx.adb_client, "is_online", True) is False: + return False + self.ctx.get_adb_client() + return True + except Exception: + return False + + @trace(type="agent", name="diagnoser") + async def run(self, prompt: str, state: State) -> str: + """Runs the Diagnoser and returns its summary.""" + + # 1. Load prompt + prompt_path = Path(__file__).parent.joinpath("diagnoser.json") + if not prompt_path.exists(): + raise FileNotFoundError(f"Prompt file not found at {prompt_path}") + + with open(prompt_path, encoding="utf-8") as f: + prompt_data = json.load(f) + + system_prompt = prompt_data.get("system_prompt", "") + if isinstance(system_prompt, list): + system_prompt = "\n".join(system_prompt) + checker_feedback_template = prompt_data.get("checker_feedback_template", "") + if isinstance(checker_feedback_template, list): + checker_feedback_template = "\n".join(checker_feedback_template) + + # 2. Prepare LLM + llm = get_llm(ctx=self.ctx, name="diagnoser") + + log_tool = get_analyze_logs_tool(self.ctx) + + read_note_tool = get_read_note_tool(self.ctx) + list_notes_tool = get_list_notes_tool(self.ctx) + adb_command_tool = get_run_short_adb_command_tool(self.ctx) + wait_tool = get_wait_tool(self.ctx) + ui_hierarchy_tool = get_ui_hierarchy_tool(self.ctx) + submit_answer_tool = get_submit_answer_tool(self.ctx) + + all_tools = [ + log_tool, + read_note_tool, + list_notes_tool, + adb_command_tool, + wait_tool, + ui_hierarchy_tool, + submit_answer_tool, + ] + if not self.is_device_online: + logger.info( + "Diagnoser running in OFFLINE environment: stripping adb short" + " command and wait tools." + ) + all_tools = [ + t + for t in all_tools + if t.name not in ("run_adb_command", "wait") + and not t.name.endswith(":run_adb_command") + and not t.name.endswith(":wait") + ] + + if self.ctx.execution_setup and self.ctx.execution_setup.video_recording_tools_enabled: + video_tool = get_video_analyzer_tool(self.ctx, role="diagnoser") + all_tools.append(video_tool) + + traced_tools = [trace_langchain_tool(t, self.ctx) for t in all_tools] + + # Bind custom tools + base_llm = llm + + # 4. Get context from DataEngine and State + initial_goal = state.initial_goal + + task_plan = "No task plan yet." + steps = [] + subgoal_hash = "default" + + if self.ctx.data_engine: + notes_dir = Path(self.ctx.data_engine.base_dir) / "notes" + task_plan_path = notes_dir / "task_plan.md" + + if task_plan_path.exists(): + try: + task_plan = task_plan_path.read_text(encoding="utf-8") + + subgoal_hash, _ = get_active_subgoal_hashes(task_plan) + except Exception as e: + logger.error(f"Failed to parse active subgoal in Diagnoser: {e}") + task_plan = f"Error reading task plan: {e}" + + steps = self.ctx.data_engine.get_agent_friendly_steps() + + # Build task tree + plan_and_history = "No plan or history available." + if self.ctx.data_engine and task_plan != "No task plan yet.": + try: + keep_hashes = get_recent_subgoal_hashes( + steps, + subgoal_hash, + window_steps=self.ctx.agent_config.history_window_steps, + ) + + plan_and_history = build_plan_and_history( + task_plan, + steps, + subgoal_hash, + keep_subgoal_hashes=keep_hashes, + ) + except Exception as e: + logger.error(f"Failed to build plan and history in Diagnoser: {e}") + plan_and_history = f"Error building plan and history: {e}" + + # Get Checker Feedback + checker_feedback = "" + if self.ctx.data_engine: + notes_dir = Path(self.ctx.data_engine.base_dir) / "notes" + verification_chat_path = notes_dir / f"verification_chat_{subgoal_hash}.json" + if verification_chat_path.exists(): + try: + turns = json.loads(verification_chat_path.read_text(encoding="utf-8")) + dialogue_lines = [] + for t in turns: + role = "Operator" if t["role"] == "operator" else "Checker" + dialogue_lines.append(f"**{role} (Round {t['round']})**:\n{t['content']}") + checker_feedback = "\n\n".join(dialogue_lines) + except Exception as e: + logger.error(f"Error reading verification chat in Diagnoser: {e}") + + # Get Current Screenshot and UI Hierarchy + latest_screenshot_b64 = None + minimal_list = "" + + if state.latest_screenshot and os.path.exists(state.latest_screenshot): + try: + with open(state.latest_screenshot, "rb") as f: + latest_screenshot_b64 = base64.b64encode(f.read()).decode("utf-8") + except Exception as e: + logger.error(f"Failed to read latest screenshot in Diagnoser: {e}") + + if state.latest_ui_hierarchy: + minimal_list = self._format_minimal_list(state.latest_ui_hierarchy) + + # Construct rich human message + content = [] + content.append({"type": "text", "text": f"Initial Goal: {initial_goal}\n"}) + content.append( + { + "type": "text", + "text": f"--- Plan & History ---\n{plan_and_history}\n", + } + ) + + if checker_feedback and checker_feedback_template: + formatted_feedback = checker_feedback_template.format(checker_feedback=checker_feedback) + content.append({"type": "text", "text": formatted_feedback}) + + if minimal_list: + content.append( + { + "type": "text", + "text": f"--- Visible UI Elements ---\n{minimal_list}\n", + } + ) + + if latest_screenshot_b64: + content.append({"type": "text", "text": "--- Current Screenshot ---"}) + content.append( + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{latest_screenshot_b64}"}, + } + ) + + content.append({"type": "text", "text": f"--- Diagnostic Query ---\n{prompt}"}) + + current_messages = [ + SystemMessage(content=system_prompt), + HumanMessage(content=content), + ] + + self.ctx.background_jobs = {} + + max_iterations = 30 + outcome = "No result" + + try: + for i in range(max_iterations): + # 5. Agent Loop Injection & Final Turn Guard + if i == max_iterations - 1: + # Identify all jobs in the registry with status "running" + running_jobs = [ + job + for job in self.ctx.background_jobs.values() + if job.get("status") == "running" + ] + if running_jobs: + # Await all associated task objects + await asyncio.gather( + *(job["task"] for job in running_jobs if job.get("task") is not None), + return_exceptions=True, + ) + + # Iterate through the background_jobs registry for all jobs where status is "completed" or "failed" and consumed is false + new_tool_messages = [] + for job in self.ctx.background_jobs.values(): + if job.get("status") in ( + "completed", + "failed", + ) and not job.get("consumed", False): + tm_status = "success" if job["status"] == "completed" else "error" + tm = ToolMessage( + tool_call_id=job["tool_call_id"], + content=job["result"] or "", + status=tm_status, + ) + new_tool_messages.append(tm) + job["consumed"] = True + + if new_tool_messages: + current_messages.extend(new_tool_messages) + + if i == max_iterations - 1: + current_messages.append( + HumanMessage( + content=( + "[WARNING] This is your final iteration. Call" + " 'submit_answer' to provide your final" + " diagnosis." + ) + ) + ) + submit_tool = next( + (t for t in traced_tools if t.name == "submit_answer"), + None, + ) + llm = base_llm.bind_tools(tools=[submit_tool] if submit_tool else []) + else: + if i > 0: + current_messages.append( + HumanMessage( + content=( + "You have not completed the diagnosis yet" + f" (iteration {i + 1} of {max_iterations})." + ) + ) + ) + llm = base_llm.bind_tools(tools=traced_tools) + + async def run_stream(): + full_response = None + trace_id = CURRENT_TRACE_ID.get() + async for chunk in llm.astream(current_messages): + if full_response is None: + full_response = chunk + else: + full_response += chunk + + pass + return full_response + + with TraceSpan(name="gemini_diagnoser_call", ctx=self.ctx) as span: + span.payload = { + "context_components": { + "system_prompt": system_prompt, + "plan_and_history": plan_and_history, + "checker_feedback": checker_feedback, + "visible_ui_elements": minimal_list, + "diagnostic_query": prompt, + } + } + response = await invoke_llm_with_timeout_message(run_stream()) + span.result = ( + response.content if hasattr(response, "content") else str(response) + ) + + if hasattr(response, "usage_metadata") and response.usage_metadata is not None: + usage = response.usage_metadata + if isinstance(usage, dict): + prompt_token_count = usage.get("prompt_token_count") + candidates_token_count = usage.get("candidates_token_count") + cached_content_token_count = usage.get("cached_content_token_count") + else: + prompt_token_count = getattr(usage, "prompt_token_count", None) + candidates_token_count = getattr(usage, "candidates_token_count", None) + cached_content_token_count = getattr( + usage, "cached_content_token_count", None + ) + + if span.payload is None: + span.payload = {} + + span.payload["token_usage"] = { + "prompt_token_count": prompt_token_count, + "candidates_token_count": candidates_token_count, + "cached_content_token_count": (cached_content_token_count), + } + + if response.content: + outcome = response.content + + if not response.tool_calls: + break + + submit_call = next( + ( + tc + for tc in response.tool_calls + if tc["name"] == "submit_answer" or tc["name"].endswith(":submit_answer") + ), + None, + ) + if submit_call: + args = submit_call["args"] + analysis = args.get("analysis", "No analysis provided.") + steps = args.get("actionable_steps", []) + + submit_tool = next( + ( + t + for t in traced_tools + if t.name == "submit_answer" or t.name.endswith(":submit_answer") + ), + None, + ) + if submit_tool: + submit_tool.invoke(args) + + outcome = f"Analysis: {analysis} | Actionalble Steps: " + for step in steps: + outcome += f"{step} " + break + + if i == max_iterations - 1: + tool_names = ", ".join(tc["name"] for tc in response.tool_calls) + warning = ( + "\n\n[Warning: Diagnoser reached the maximum iteration" + f" limit of {max_iterations} and could not execute" + f" tool(s): {tool_names}]" + ) + if outcome == "No result": + outcome = ( + "Diagnoser reached the maximum iteration limit of" + f" {max_iterations} while attempting to call tools:" + f" {tool_names}" + ) + else: + outcome += warning + break + + current_messages.append(response) + + async def run_tool(tc): + tool_name = tc["name"].split(":")[-1] if ":" in tc["name"] else tc["name"] + logger.info(f"Diagnoser requested tool: {tool_name}") + + result = "" + status = "success" + try: + if tool_name == "video_analyzer": + job_id = f"job_{uuid.uuid4().hex}" + self.ctx.background_jobs[job_id] = { + "tool_call_id": tc["id"], + "status": "running", + "result": None, + "consumed": False, + "task": None, + } + if self.ctx.data_engine: + self.ctx.data_engine.register_background_task( + task_id=job_id, + summary=(f"Diagnoser Video Analyzer Tool Call: {tc['id']}"), + ) + + async def run_bg_job(): + try: + tool_to_run = next( + (t for t in traced_tools if t.name == "video_analyzer"), + None, + ) + if tool_to_run: + args = dict(tc["args"]) + result_obj = await invoke_tool_with_injection( + tool=tool_to_run, + args=args, + tool_call_id=tc["id"], + state=state, + ) + result_content = get_tool_result_content(result_obj) + self.ctx.background_jobs[job_id]["result"] = result_content + self.ctx.background_jobs[job_id]["status"] = "completed" + if self.ctx.data_engine: + self.ctx.data_engine.unregister_background_task( + task_id=job_id, + status="completed", + ) + else: + self.ctx.background_jobs[job_id]["result"] = ( + "Error: video_analyzer tool not found" + ) + self.ctx.background_jobs[job_id]["status"] = "failed" + if self.ctx.data_engine: + self.ctx.data_engine.unregister_background_task( + task_id=job_id, status="failed" + ) + except asyncio.CancelledError: + logger.info(f"Background job {job_id} cancelled.") + raise + except Exception as e: + logger.error(f"Background job {job_id} failed: {e}") + self.ctx.background_jobs[job_id]["result"] = str(e) + self.ctx.background_jobs[job_id]["status"] = "failed" + if self.ctx.data_engine: + self.ctx.data_engine.unregister_background_task( + task_id=job_id, status="failed" + ) + + task = asyncio.create_task(run_bg_job()) + self.ctx.background_jobs[job_id]["task"] = task + + return ToolMessage( + tool_call_id=tc["id"], + content=( + f"Video analyzer started. Job ID: {job_id}." + " [Warning] No background logs, status" + " updates, or errors will be logged." + ), + status="success", + ) + + if ":" in tool_name: + tool_to_run = get_tool_by_name(tool_name, traced_tools) + else: + tool_to_run = next( + (t for t in traced_tools if t.name == tool_name), + None, + ) + if tool_to_run: + args = dict(tc["args"]) + with TraceSpan(name=tool_name, ctx=self.ctx) as span: + span.payload = {"args": args} + try: + result_obj = await invoke_tool_with_injection( + tool=tool_to_run, + args=args, + tool_call_id=tc["id"], + state=state, + ) + result = get_tool_result_content(result_obj) + if isinstance(result, list): + result_str = "\n".join(map(str, result)) + elif not isinstance(result, str): + result_str = str(result) + else: + result_str = result + span.result = result + if result_str.startswith("Error"): + status = "error" + span.status = "failed" + span.error = result + except Exception as e: + status = "error" + span.status = "failed" + span.error = str(e) + raise e + else: + result = f"Error: Tool {tool_name} not supported" + status = "error" + except Exception as e: + logger.error(f"Error running tool {tool_name}: {e}") + result = f"Error running tool {tool_name}: {e}" + status = "error" + + return ToolMessage( + tool_call_id=tc["id"], + content=result, + status=status, + ) + + # Execute all requested tools in parallel + active_tool_calls = [ + tc + for tc in response.tool_calls + if (tc["name"].split(":")[-1] if ":" in tc["name"] else tc["name"]) + != "google_search" + ] + tool_outputs = await asyncio.gather(*(run_tool(tc) for tc in active_tool_calls)) + + for tm in tool_outputs: + current_messages.append(tm) + finally: + # 6. Cleanup on Exit + for job in self.ctx.background_jobs.values(): + if job.get("status") == "running" and job.get("task") is not None: + job["task"].cancel() + + return outcome + + def _format_minimal_list(self, fused_xml: list[dict[str, Any]]) -> str: + """Formats fused XML into a minimal structured list.""" + lines = [] + for i, node in enumerate(fused_xml): + text = node.get("text") or node.get("content-desc") or "" + bounds = node.get("bounds") + + ocr_elements = node.get("ocr_elements") + if ocr_elements: + for ocr in ocr_elements: + lines.append(f"[{i}] OCR Text: '{ocr['text']}' | Bounds: {ocr['bounds']}") + elif text.strip(): + lines.append(f"[{i}] Text: '{text.strip()}' | Bounds: {bounds}") + + return "\n".join(lines) diff --git a/artemis/agents/explorer/constants.py b/artemis/agents/explorer/constants.py new file mode 100644 index 0000000..dea9115 --- /dev/null +++ b/artemis/agents/explorer/constants.py @@ -0,0 +1,72 @@ +# 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. + +"""Constants and prompt templates for the UI Explorer agent.""" + +EXPLORE_DESCRIPTIONS = { + "flash": { + "description": ( + "[EXPLORER] Call this to ask the UI Explorer agent to " + "locate coordinates on the screen layout by object detection." + ), + "query_description": ( + "Pipe-separated list of target element descriptions, semantic" + " action terms, or color descriptors to search for concurrently" + " (e.g., 'red settings button | gear icon')." + ), + "rule_prompt": ( + "Use the explorer tool to quickly check for visual icons, buttons, or text by name." + ), + }, + "pro": { + "description": ( + "[EXPLORER] Call this to ask the UI Explorer agent to locate" + " coordinates on the screen layout by executing XML, OCR," + " coordinate search, and object detection." + ), + "query_description": ( + "The target element or information to search for, including descriptions." + ), + "rule_prompt": ( + "Call the explorer tool to search, OCR, or inspect it (maximum {max_iterations} tries)." + ), + "version_prompt": ( + "You are running with a maximum of {max_iterations} turns: on your" + " final turn, you will receive a warning reminding you to call" + " submit_answer." + ), + }, + "ultra": { + "description": ( + "[EXPLORER] Full visual reasoning agent. Deeply inspects screen" + " regions and utilizes all perception/pixel-level tools with a high" + " execution limit (up to {max_iterations} turns) for" + " high-difficulty or layout-critical tasks." + ), + "query_description": ( + "The target element or information to search for, including descriptions." + ), + "rule_prompt": ( + "For highly complex visual searches, details validation, or layout" + " verification, invoke the explorer tool to perform multi-turn" + " visual reasoning and detailed coordinate checks (maximum" + " {max_iterations} tries)." + ), + "version_prompt": ( + "You are running with a maximum of {max_iterations} turns: on your" + " final turn, you will receive a warning reminding you to call" + " submit_answer." + ), + }, +} diff --git a/artemis/agents/explorer/explorer.json b/artemis/agents/explorer/explorer.json new file mode 100644 index 0000000..914e3d5 --- /dev/null +++ b/artemis/agents/explorer/explorer.json @@ -0,0 +1,27 @@ +{ + "IDENTITY": "You are the agentic UI Explorer designed to identify objects, scan text, and extract UI attributes on a phone screen.", + "OPERATING PRINCIPLES": [ + "**Confidence-Driven Cost-Benefit Strategy**: Maximize localization accuracy while minimizing execution latency and tool calls:", + " - **Direct Submit (High Confidence)**: If you are highly confident in the target coordinates from the initial annotated screenshot and minimal list, call `submit_answer` immediately in Turn 1 to save time.", + " - **Cost-Benefit Turn Planning (Low Confidence / Ambiguity)**: If you are not highly confident. Consider structuring your exploration turns by weighing the latency/cost of each tool against the expected accuracy payoff:", + " - **Parallel Context Gathering**: MUST parallel calling all perception tools in the first turn to gather broad screen context.", + " - **Targeted Verification**: If micro-details or coordinates are still ambiguous, zoom in on promising candidate areas (parallel tool calls). Note that a single zoom-in or multi-zoom-in operation takes one turn, plan your verification steps accordingly.", + " - **Finalization**: Submit coordinates using `submit_answer`. Stop exploring as soon as further tool calls would yield negligible accuracy gains relative to their latency.", + "**Tool Capabilities**: You have three classes of tools: `Perception` tools (to extract and locate elements, coordinates, or inspect specific regions), `Transformation` tools (for pixel-level image editing), and the `Finalization` tool (to submit final answers and terminate). Note that in your final turn, all other tools are stripped away, leaving only the finalization tool, forcing you to submit your final answer.", + "**Parallel vs Sequential Invocations**: You are fully supported to call perception tools in parallel (within a single turn) to expedite your search. However, transformation tools are expected to be called sequentially (since they transform the image for subsequent perception steps).", + "**Multimodal Perception**: You have access to images marked with bounding boxes and labels. Use both the visual information in the images and the structured text lists to make decisions.", + "**Image Pool Tracking**: Images are managed in an Image Pool (starting with 'img_0'). Transformations via transformation tools will register new image IDs (e.g., 'img_1').", + "**Prefix System**: Elements in the screenshots are labeled with prefixes returned by your tools (e.g., Numbers for initial UI, letters for tool outputs). Always reference these precise labels (e.g., \"I choose S2\") when making decisions.", + "**Contextual Awareness**: Always prioritize checking the context feedback to learn from previous mistakes.", + "**Ask Perception Tool**: Call the unified `ask_perception_tool` specifying `search_query`, `nx`, `ny`, and `detect_queries` to perform text lookups, coordinate queries, and visual detection concurrently.", + "**Locating Components**: Locate target components on screen by matching the X-labels (from text/coordinate searches) and O-labels (from object detection) in the text response with their corresponding green, blue, or red indicator dots on the output images." + ], + "BEHAVIORAL GUIDELINES & SOFT CONSTRAINTS": [ + "**Avoid Hallucination**: Avoid guessing coordinates. If you are unable to locate the target element on the screen after trying all perception options, submit `[]` in `candidates` and document your observations as specified in the `submit_answer` schema.", + "**Ambiguity Resolution**: If the screen contains multiple identical elements or several candidates that match the query (e.g., three \"Add\" buttons), you are encouraged to return a ranked list of all viable candidates (up to 10 candidates) via `submit_answer`, ordered by your confidence, with brief descriptions to distinguish them.", + "**Semantic & Icon Synonym Fallback**: Literal exact matches are ideal, but semantic equivalents and synonyms are fully acceptable. If a text element is requested but only a corresponding icon is available, you should attempt to map the synonym and locate the icon (e.g., map \"Search\" to a magnifying glass icon, map \"Add\" to \"+\" symbol)." + ], + "STRICT CONSTRAINTS": [ + "**Validation Standards**: All submitted coordinates `[nx, ny]` for candidates in `submit_answer` must strictly be integers in the `[0-1000]` range inclusive. If your parameters fail business validation rules, you will receive dynamic error responses in your execution loop, allowing you to self-correct and re-submit." + ] +} \ No newline at end of file diff --git a/artemis/agents/explorer/explorer.py b/artemis/agents/explorer/explorer.py new file mode 100644 index 0000000..4464466 --- /dev/null +++ b/artemis/agents/explorer/explorer.py @@ -0,0 +1,2260 @@ +# 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 asyncio +import base64 +import glob +import hashlib +import json +import os +from pathlib import Path +import random +import re +import time +from typing import Any, Literal + +import cv2 +from google import genai +from google.genai import types +from google.genai.errors import APIError +import httpx +from langchain_core.messages import ( + BaseMessage, + HumanMessage, + SystemMessage, + ToolMessage, +) + +from artemis.agents.image_processor.image_processor import ImageProcessor +from artemis.agents.object_detector.object_detector import _run_object_detection +from artemis.config import settings +from artemis.constants import SAFETY_SETTINGS_BLOCK_NONE +from artemis.context import ArtemisContext +from artemis.data_engine.storage import StorageManager +from artemis.data_engine.trace import TraceSpan, trace +from artemis.graph.state import State +from artemis.services.llm import get_llm + +# Import diagnostic functions directly +from artemis.agents.explorer.constants import EXPLORE_DESCRIPTIONS +from artemis.mcp.xml_search_server import ( + search_by_coordinates as search_by_coordinates_func, + search_ui as search_ui_func, +) +from artemis.utils.logger import get_logger +from artemis.utils.ocr_api import perform_ocr +from artemis.utils.ocr_xml_fusion import fuse_ocr_with_xml +from artemis.utils.visualization import draw_dots, format_minimal_list_with_points + +logger = get_logger(__name__) + +UNIVERSAL_EXPLORER_TOOLS = [ + { + "type": "function", + "function": { + "name": "ask_perception_tool", + "description": ( + "[Perception] Concurrently executes and awaits three sub-tasks in parallel," + " including database search, coordinate search, and visual object detection." + ), + "parameters": { + "type": "object", + "properties": { + "search_query": { + "type": "string", + "description": "Required. Text query to search for in XML/OCR elements.", + }, + "nx": { + "type": "integer", + "description": "Required. Normalized X coordinate (0-1000 range).", + }, + "ny": { + "type": "integer", + "description": "Required. Normalized Y coordinate (0-1000 range).", + }, + "detect_queries": { + "type": "array", + "items": {"type": "string"}, + "description": "Required. List of query strings to visually detect.", + }, + }, + "required": ["search_query", "nx", "ny", "detect_queries"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "detect_objects", + "description": ( + "[Perception] Locates elements, shapes and more using powerful VLM." + " Returns their normalized coordinates (0-1000 scale) with label dots." + ), + "parameters": { + "type": "object", + "properties": { + "target_image_id": { + "type": "string", + "description": "The ID of the image to process.", + }, + "queries": { + "type": "array", + "items": {"type": "string"}, + "description": "A list of single query strings to detect on the screen.", + }, + }, + "required": ["queries", "target_image_id"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "get_ocr_list", + "description": ( + "[Perception] Retrieves all text elements detected on the screen via OCR." + ), + "parameters": {"type": "object", "properties": {}}, + }, + }, + { + "type": "function", + "function": { + "name": "ask_image_processor", + "description": ("[Transformation] A scripting tool for pixel-level image processing."), + "parameters": { + "type": "object", + "properties": { + "target_image_id": { + "type": "string", + "description": "The ID of the image in the Image Pool to start from.", + }, + "instruction": { + "type": "string", + "description": "1-2 precise commands describing a simple task.", + }, + }, + "required": ["instruction", "target_image_id"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "inspect_region", + "description": ( + "[Perception] Crops and resizes a bounding box region of the original screenshot." + ), + "parameters": { + "type": "object", + "properties": { + "x_min": {"type": "integer", "description": "Left X coordinate (0-1000)."}, + "y_min": {"type": "integer", "description": "Top Y coordinate (0-1000)."}, + "x_max": {"type": "integer", "description": "Right X coordinate (0-1000)."}, + "y_max": {"type": "integer", "description": "Bottom Y coordinate (0-1000)."}, + "zoom_factor": { + "type": "number", + "description": "Zoom factor to upscale the cropped region (1.0 to 4.0).", + }, + }, + "required": ["x_min", "y_min", "x_max", "y_max", "zoom_factor"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "submit_answer", + "description": ( + "[Finalization] Submits the final list of ranked candidate UI elements matching" + " the search query, or a fallback message if no elements were found." + ), + "parameters": { + "type": "object", + "properties": { + "candidates": { + "type": "array", + "items": { + "type": "object", + "properties": { + "label": { + "type": "string", + "description": "Reference label (e.g. S1, D2, T3)", + }, + "description": { + "type": "string", + "description": "Brief description of this candidate element", + }, + "coords": { + "type": "array", + "items": {"type": "integer"}, + "description": ( + "[nx, ny] coordinates in normalized 0-1000 scale" + ), + }, + }, + "required": ["label", "coords"], + }, + "description": "A ranked list of candidate UI elements matching the query.", + }, + "fallback_message": { + "type": "string", + "description": ( + "Optional fallback explanation if target elements were not found." + ), + }, + }, + "required": ["candidates"], + }, + }, + }, +] + + +class Explorer: + BLACKLIST_TEMPLATE = ( + "\n# TOOL BLACKLIST\n- The following tools are blacklisted and cannot be used: {tools}\n" + ) + TOOLS = [ + types.FunctionDeclaration( + name="ask_perception_tool", + description=( + "[Perception] Concurrently executes and awaits three sub-tasks" + " in parallel, including database search, coordinate search," + " and visual object detection on the screen.\n\nOutput" + " Format:\nReturns:'text' (str): A single string consolidating" + " results with 0-1000 normalized coordinates, including xml" + " search matches labeled [X], ocr search matches labeled [O]," + " coordinate search matches labeled [X], and visually detected" + " object targets labeled [D].\n(Annotated screenshots are also" + " added to the context containing corresponding colored label" + " dots)." + ), + parameters=types.Schema( + type=types.Type.OBJECT, + properties={ + "search_query": types.Schema( + type=types.Type.STRING, + description=("Required. Text query to search for in XML/OCR elements."), + ), + "nx": types.Schema( + type=types.Type.INTEGER, + description=( + "Required. Normalized X coordinate (0-1000 range)" + " to audit for overlapping elements." + ), + ), + "ny": types.Schema( + type=types.Type.INTEGER, + description=( + "Required. Normalized Y coordinate (0-1000 range)" + " to audit for overlapping elements." + ), + ), + "detect_queries": types.Schema( + type=types.Type.ARRAY, + items=types.Schema(type=types.Type.STRING), + description=( + "Required. List of query strings to visually detect on the screen." + ), + ), + }, + required=["search_query", "nx", "ny", "detect_queries"], + ), + ), + types.FunctionDeclaration( + name="detect_objects", + description=( + "[Perception] Locates elements, shapes and more using powerful" + " VLM. Returns their normalized coordinates (in a [0, 1000]" + " scale) labeled with [D1], [D2] etc. in the text output, along" + " with an annotated image containing corresponding visual label" + " dots." + ), + parameters=types.Schema( + type=types.Type.OBJECT, + properties={ + "target_image_id": types.Schema( + type=types.Type.STRING, + description="The ID of the image to process.", + ), + "queries": types.Schema( + type=types.Type.ARRAY, + items=types.Schema(type=types.Type.STRING), + description=("A list of single query strings to detect on the screen."), + ), + }, + required=["queries", "target_image_id"], + ), + ), + types.FunctionDeclaration( + name="get_ocr_list", + description=( + "[Perception] Retrieves all text elements detected on the" + " original, unprocessed screen via OCR. Returns their" + " normalized coordinates (in a [0, 1000] scale) labeled with" + " [T1], [T2] etc. in the text output, along with an annotated" + " image containing corresponding visual label dots." + ), + parameters=types.Schema( + type=types.Type.OBJECT, + properties={}, + ), + ), + types.FunctionDeclaration( + name="ask_image_processor", + description=( + "[Transformation] A scripting tool for pixel-level image" + " processing. CAUTION: This tool is slow and expensive. Returns" + " a summary of the operations performed and new image IDs," + " along with new labels and screen coordinates." + ), + parameters=types.Schema( + type=types.Type.OBJECT, + properties={ + "target_image_id": types.Schema( + type=types.Type.STRING, + description=( + "The ID of the image in the Image Pool to start" + " from (e.g. 'img_0', 'img_1')." + ), + ), + "instruction": types.Schema( + type=types.Type.STRING, + description=( + "1-2 precise commands describing a simple task." + " Include all necessary coordinates and other" + " parameters.\n" + ), + ), + }, + required=["instruction", "target_image_id"], + ), + ), + types.FunctionDeclaration( + name="submit_answer", + description=( + "[Finalization] Submits the final list of ranked candidate UI" + " elements matching the search query, or a fallback message if" + " no elements were found, to complete the search task. Fails" + " validation if: (1) candidates is empty and no" + " fallback_message is provided, (2) coordinates are not exactly" + " 2 integers [nx, ny], or (3) coordinates lie outside" + " normalized [0, 1000] range." + ), + parameters=types.Schema( + type=types.Type.OBJECT, + properties={ + "candidates": types.Schema( + type=types.Type.ARRAY, + items=types.Schema( + type=types.Type.OBJECT, + properties={ + "label": types.Schema( + type=types.Type.STRING, + description=( + "Reference label (e.g., S1, D2, T3, or number like 3)" + ), + ), + "coords": types.Schema( + type=types.Type.ARRAY, + items=types.Schema(type=types.Type.INTEGER), + description=("[nx, ny] coordinates in normalized 0-1000 scale"), + ), + "description": types.Schema( + type=types.Type.STRING, + description=("Brief description of this candidate element"), + ), + }, + required=["label", "coords"], + ), + description=( + "A ranked list of candidate UI elements that match" + " the query, sorted from highest to lowest" + " confidence (maximum 10 candidates). Leave empty" + " if no matching elements were found on the screen." + ), + ), + "fallback_message": types.Schema( + type=types.Type.STRING, + description=( + "Optional fallback explanation if target elements" + " were not found or if there are special" + " observations." + ), + ), + }, + required=["candidates"], + ), + ), + types.FunctionDeclaration( + name="inspect_region", + description=( + "[Perception] Crops and resizes a bounding box region of the" + " original screenshot. Returns the cropped and resized image" + " directly, without any image ID or coordinate transform." + ), + parameters=types.Schema( + type=types.Type.OBJECT, + properties={ + "x_min": types.Schema( + type=types.Type.INTEGER, + description="Left X coordinate (0-1000).", + ), + "y_min": types.Schema( + type=types.Type.INTEGER, + description="Top Y coordinate (0-1000).", + ), + "x_max": types.Schema( + type=types.Type.INTEGER, + description="Right X coordinate (0-1000).", + ), + "y_max": types.Schema( + type=types.Type.INTEGER, + description="Bottom Y coordinate (0-1000).", + ), + "zoom_factor": types.Schema( + type=types.Type.NUMBER, + description=("Zoom factor to upscale the cropped region (1.0 to 4.0)."), + ), + }, + required=["x_min", "y_min", "x_max", "y_max", "zoom_factor"], + ), + ), + ] + + def __init__(self, ctx: ArtemisContext): + self.ctx = ctx + self.global_label_idx = 1 + self.width = 1080 + self.height = 2400 + self.image_name = None + self.screenshot_path = None + self.image_pool = {} + self.next_img_id = 1 + try: + agent_cfg = getattr(self.ctx, "agent_config", None) + blacklisted_config = ( + getattr(agent_cfg, "blacklisted_tools", {}).get("explorer", []) if agent_cfg else [] + ) + self.blacklisted_tools = set(blacklisted_config) + except (TypeError, AttributeError): + self.blacklisted_tools = set() + self.http_client = None + self.turn_latencies = [] + self.turn_cached_tokens = [] + self.trace_history = [] + self._init_engine() + + def _init_engine(self) -> None: + """Initializes model engine and decides whether to use native Gemini or Universal path.""" + ctx = self.ctx + llm_config = getattr(ctx, "llm_config", None) + llm_cfg = getattr(llm_config, "explorer", None) if llm_config else None + model_str = ( + llm_cfg.model if (llm_cfg and hasattr(llm_cfg, "model")) else "gemini-3.6-flash" + ).lower() + self.model_name = ( + llm_cfg.model if (llm_cfg and hasattr(llm_cfg, "model")) else "gemini-3.6-flash" + ) + if "/" in self.model_name: + self.model_name = self.model_name.split("/")[-1] + + has_google_key = bool( + settings.GOOGLE_API_KEY and settings.GOOGLE_API_KEY.get_secret_value() + ) + is_gemini_model = "gemini" in model_str + + self.client = getattr(ctx, "_genai_client", None) + if self.client is not None: + self.use_native_gemini = True + elif has_google_key and is_gemini_model: + try: + self.client = genai.Client(api_key=settings.GOOGLE_API_KEY.get_secret_value()) + ctx._genai_client = self.client + self.use_native_gemini = True + except Exception as e: + logger.warning( + f"Failed to initialize native Gemini client for Explorer: {e}." + " Using universal engine." + ) + self.use_native_gemini = False + else: + self.use_native_gemini = False + + def _prune_historical_images(self, contents, keep_last=1): + image_parts = [] + for content in contents: + for part in content.parts: + if getattr(part, "inline_data", None) or getattr(part, "file_data", None): + image_parts.append(part) + if len(image_parts) <= keep_last: + return + to_keep_ids = {id(p) for p in image_parts[-keep_last:]} + for content in contents: + content.parts = [ + types.Part.from_text( + text=("[Image pruned to maintain visual focus on latest state]") + ) + if (getattr(p, "inline_data", None) or getattr(p, "file_data", None)) + and id(p) not in to_keep_ids + else p + for p in content.parts + ] + + def get_exposed_tools(self, only_submit: bool = False) -> list[types.FunctionDeclaration]: + """Returns the list of tools, filtering out blacklisted ones.""" + tools = [tool for tool in self.TOOLS if tool.name not in self.blacklisted_tools] + if only_submit: + tools = [tool for tool in tools if tool.name == "submit_answer"] + return tools + + async def _search_ui_helper(self, query: str, prefix: str = "S", color: str = "red") -> dict: + if not self.image_name: + return { + "text": ( + "Error: UI data for the current screen is not found in Data" + " Engine. It might not be synced yet." + ), + "image_path": None, + } + + # 1. Try with high precision threshold + raw_res = search_ui_func(self.image_name, query, 0.7) + if "error" in raw_res: + return {"text": raw_res["error"], "image_path": None} + + matches = raw_res.get("matches", []) + + # 2. If fewer than 3 matches, try with a lower threshold and merge + if len(matches) < 3: + raw_res_low = search_ui_func(self.image_name, query, 0.4) + if "error" not in raw_res_low: + low_matches = raw_res_low.get("matches", []) + existing_keys = { + ( + m.get("matched_text"), + tuple(m.get("bounds")) if m.get("bounds") else None, + ) + for m in matches + } + for m in low_matches: + key = ( + m.get("matched_text"), + tuple(m.get("bounds")) if m.get("bounds") else None, + ) + if key not in existing_keys: + matches.append(m) + existing_keys.add(key) + + if not matches: + return {"text": "No matches found.", "image_path": None} + + points = [] + labels = [] + text_output = [] + + for m in matches: + m_type = m.get("type", "xml") + m_prefix = "O" if m_type == "ocr" else prefix + label_id = f"{m_prefix}{self.global_label_idx}" + self.global_label_idx += 1 + bounds = m.get("bounds") + matched_text = m.get("matched_text", "") + + if bounds and len(bounds) == 4: + left, top, right, bottom = bounds + cx_pixel = (left + right) // 2 + cy_pixel = (top + bottom) // 2 + + points.append([cx_pixel, cy_pixel]) + labels.append(label_id) + + cx_norm = int(max(0, min(1000, cx_pixel * 1000 / self.width))) + cy_norm = int(max(0, min(1000, cy_pixel * 1000 / self.height))) + + text_output.append(f"[{label_id}] '{matched_text}' at [{cx_norm},{cy_norm}]") + else: + text_output.append(f"[{label_id}] '{matched_text}' at unknown") + + if not points: + return {"text": " | ".join(text_output), "image_path": None} + + base_dir = settings.TRACES_PATH + images_dir = base_dir / "images" + search_ui_dir = images_dir / "search_ui" + search_ui_dir.mkdir(parents=True, exist_ok=True) + + existing_files = glob.glob(str(search_ui_dir / f"{self.image_name}_*.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 + output_path = search_ui_dir / f"{self.image_name}_{seq}.jpg" + + try: + draw_dots( + self.screenshot_path, + points, + labels, + str(output_path), + color=color, + ) + image_path = str(output_path) + except Exception as e: + logger.warning(f"Failed to draw dots: {e}") + image_path = None + + return {"text": " | ".join(text_output), "image_path": image_path} + + async def _search_by_coords_helper( + self, nx: int, ny: int, prefix: str = "X", color: str = "blue" + ) -> dict: + if not self.image_name: + return { + "text": ( + "Error: UI data for the current screen is not found in Data" + " Engine. It might not be synced yet." + ), + "image_path": None, + } + + x = int(max(0, min(self.width, nx * self.width / 1000))) + y = int(max(0, min(self.height, ny * self.height / 1000))) + raw_res = search_by_coordinates_func(self.image_name, x, y) + + label_id = f"{prefix}{self.global_label_idx}" + self.global_label_idx += 1 + + # Draw a dot at the query coordinates + base_dir = settings.TRACES_PATH + images_dir = base_dir / "images" + coords_audit_dir = images_dir / "coords_audit" + coords_audit_dir.mkdir(parents=True, exist_ok=True) + + existing_files = glob.glob(str(coords_audit_dir / f"{self.image_name}_*.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 + output_path = coords_audit_dir / f"{self.image_name}_{seq}.jpg" + + try: + draw_dots( + self.screenshot_path, + [[x, y]], + [label_id], + str(output_path), + color=color, + ) + image_path = str(output_path) + except Exception as e: + logger.warning(f"Failed to draw blue dots: {e}") + image_path = None + + raw_res_str = str(raw_res).replace("\n", " | ").replace(" | | ", " | ") + return { + "text": (f"Matched element at [{nx},{ny}]: [{label_id}] | {raw_res_str}"), + "image_path": image_path, + } + + @trace(type="tool", name="detect_objects") + async def exec_detect_objects( + self, + queries: list[str], + target_image_id: str = "img_0", + prefix: str = "D", + color: str = "red", + ) -> dict: + try: + target_info = self.image_pool.get(target_image_id) + if not target_info: + return { + "text": (f"Error: {target_image_id} not found in Image Pool."), + "image_path": None, + } + + target_path = target_info["path"] + transform = target_info["transform"] + + # Load object detector templates + detector_prompt_path = ( + Path(__file__).parent.parent / "object_detector" / "object_detector.json" + ) + global_timeout = None + try: + with open(detector_prompt_path, encoding="utf-8") as f: + detector_config = json.load(f) + detector_templates = detector_config.get("templates", []) + detector_instructions = detector_config.get("instructions", "") + global_timeout = detector_config.get("global_timeout", None) + except Exception as e: + logger.warning(f"Failed to load detector prompt config: {e}") + detector_templates = [ + "Point to the following objects in the provided image: {labels_str}." + ] + detector_instructions = "" + + templates = [f"{t}\n\n{detector_instructions}" for t in detector_templates] + + result = await _run_object_detection( + self.ctx, + target_path, + queries, + templates, + global_timeout=global_timeout, + ) + try: + detected_items = result.get("detected", []) + failed_queries = result.get("failed", []) + + if not detected_items: + text_output = ["No objects detected."] + if failed_queries: + text_output.append( + "\n[NOTE: Failed to detect the following queries:" + f" {', '.join(failed_queries)}]" + ) + return { + "text": "".join(text_output), + "image_path": str(target_path), + } + + base_dir = settings.TRACES_PATH + images_dir = base_dir / "images" + object_detect_dir = images_dir / "object_detect" + object_detect_dir.mkdir(parents=True, exist_ok=True) + + image_name_safe = self.image_name or "temp_image" + existing_files = glob.glob(str(object_detect_dir / f"{image_name_safe}_*.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 + output_path = object_detect_dir / f"{image_name_safe}_{seq}.jpg" + + points = [] + d_labels = [] + text_output = [] + + if target_image_id != "img_0": + text_output.append( + "[NOTE: Coordinates are automatically mapped back to" + " the original 1080x2400 screen space from" + f" {target_image_id}]" + ) + + for item in detected_items: + label_id = f"{prefix}{self.global_label_idx}" + self.global_label_idx += 1 + pos = item.get("point") + if pos and isinstance(pos, list) and len(pos) == 2: + x_norm, y_norm = pos + + t_img = cv2.imread(target_path) + t_h, t_w = t_img.shape[:2] + + # Pixels in target image + x_target_pixel = (x_norm / 1000) * t_w + y_target_pixel = (y_norm / 1000) * t_h + + # Pixels in original image + x_orig_pixel = (x_target_pixel / transform["scale_x"]) + transform[ + "offset_x" + ] + y_orig_pixel = (y_target_pixel / transform["scale_y"]) + transform[ + "offset_y" + ] + + points.append([int(x_orig_pixel), int(y_orig_pixel)]) + d_labels.append(label_id) + + # Original image norm + x_orig_norm = int(max(0, min(1000, x_orig_pixel * 1000 / self.width))) + y_orig_norm = int(max(0, min(1000, y_orig_pixel * 1000 / self.height))) + + text_output.append( + f'[{label_id}] "{item.get("label")}" at [{x_orig_norm},{y_orig_norm}]' + ) + + draw_dots( + self.screenshot_path, + points, + d_labels, + str(output_path), + color=color, + ) + + if failed_queries: + text_output.append( + " | [NOTE: Failed to detect the following queries:" + f" {', '.join(failed_queries)}]" + ) + + return { + "text": " | ".join(text_output), + "image_path": str(output_path), + } + + except Exception as e: + logger.warning(f"Failed to process detection output for annotation: {e}") + return {"text": str(result), "image_path": None} + except Exception as e: + return { + "text": f"Error: Object detection failed: {e}", + "image_path": None, + } + + @trace(type="tool", name="ask_perception_tool") + async def exec_ask_perception_tool( + self, search_query: str, nx: int, ny: int, detect_queries: list[str] + ) -> dict: + if nx is None or ny is None or not (0 <= nx <= 1000) or not (0 <= ny <= 1000): + coords_invalid = True + else: + coords_invalid = False + + # Run tasks concurrently in parallel! + tasks = [ + self._search_ui_helper(search_query, prefix="X", color="green"), + self._search_by_coords_helper(nx, ny, prefix="X", color="blue") + if not coords_invalid + else asyncio.sleep( + 0, + result={ + "text": "Error: Invalid coordinates format", + "image_path": None, + }, + ), + self.exec_detect_objects( + detect_queries, target_image_id="img_0", prefix="D", color="red" + ), + ] + + results = await asyncio.gather(*tasks) + + # Consolidate text results + text_parts = [] + + xml_text = results[0].get("text", "No matches found.") + if xml_text and xml_text != "No matches found.": + text_parts.append(f"XML/OCR Text Search Results are: {xml_text}") + + coord_text = results[1].get("text", "No elements found.") + if coord_text and coord_text != "No elements found.": + text_parts.append(f"Coordinate Search Results are: {coord_text}") + + obj_text = results[2].get("text", "No objects detected.") + if obj_text: + if obj_text.startswith("No objects detected."): + note_part = obj_text.replace("No objects detected.", "").strip() + if note_part: + text_parts.append(f"Object Detection Results are: {note_part}") + else: + text_parts.append(f"Object Detection Results are: {obj_text}") + + if not text_parts: + unified_text = "No matches or objects found across any perception method." + else: + unified_text = ". ".join(text_parts) + "." + + # Collect image paths + image_paths = [] + for r in results: + if isinstance(r, dict) and r.get("image_path"): + image_paths.append(r["image_path"]) + + return { + "text": unified_text, + "image_paths": image_paths, + } + + @trace(type="tool", name="ask_image_processor") + async def exec_ask_image_processor( + self, instruction: str, target_image_id: str = "img_0" + ) -> dict: + coder = ImageProcessor(self.ctx) + target_info = self.image_pool.get(target_image_id) + if not target_info: + return { + "text": f"Error: {target_image_id} not found in Image Pool.", + "image_path": None, + } + + target_path = target_info["path"] + result = await coder.run(instruction, target_path) + + if not isinstance(result, dict) or "error" in result: + err_msg = result.get("error") if isinstance(result, dict) else "Unknown error" + return { + "text": f"Image Processor failed: {err_msg}", + "image_paths": [], + } + + outputs = result.get("outputs", []) + summary = result.get("summary", "") + + parent_transform = target_info["transform"] + registered_ids = [] + annotations_sections = [] + image_paths = [] + + for entry in outputs: + new_id = f"img_{self.next_img_id}" + self.next_img_id += 1 + + new_path = entry["path"] + image_paths.append(new_path) + transform = entry["transform"] + + final_scale_x = parent_transform["scale_x"] * transform["scale_x"] + final_scale_y = parent_transform["scale_y"] * transform["scale_y"] + final_offset_x = ( + transform["offset_x"] / parent_transform["scale_x"] + ) + parent_transform["offset_x"] + final_offset_y = ( + transform["offset_y"] / parent_transform["scale_y"] + ) + parent_transform["offset_y"] + + self.image_pool[new_id] = { + "path": new_path, + "transform": { + "offset_x": final_offset_x, + "offset_y": final_offset_y, + "scale_x": final_scale_x, + "scale_y": final_scale_y, + }, + "description": (f"Generated by coder from {target_image_id}: {summary}"), + } + registered_ids.append(new_id) + + # Parse and translate annotations + annotations = entry.get("annotations", {}) + if annotations: + lines = [f"{new_id}:"] + for label, coord in sorted(annotations.items(), key=lambda item: item[0]): + px, py = coord + sx = int(round((px / final_scale_x) + final_offset_x)) + sy = int(round((py / final_scale_y) + final_offset_y)) + lines.append(f" - [{label}]: [{sx}, {sy}]") + annotations_sections.append("\n".join(lines)) + + if registered_ids: + ids_str = ", ".join(registered_ids) + text = ( + "Image Processor completed successfully. New image(s) saved as" + f" {ids_str} in Image Pool. Summary: {summary}" + ) + if annotations_sections: + text += "\n\nAnnotations (wrt original screenshot):\n" + "\n".join( + annotations_sections + ) + return { + "text": text, + "image_paths": image_paths, + } + else: + return { + "text": f"ImageProcessor error: {err_msg}", + "image_paths": [], + } + + output_path = result.get("output_path") + trans = result.get("transform") + new_image_id = f"img_{len(self.image_pool)}" + if output_path and trans: + self.image_pool[new_image_id] = { + "path": output_path, + "transform": trans, + } + + text_msg = ( + f"Image processed successfully. New image saved as" + f" '{new_image_id}'. You can now use '{new_image_id}' as" + " 'target_image_id' in subsequent tool calls." + ) + return { + "text": text_msg, + "image_path": output_path, + "image_id": new_image_id, + } + + @trace(type="tool", name="get_ocr_list") + async def exec_get_ocr_list(self, prefix: str = "O", color: str = "green") -> dict: + if not self.image_name: + return { + "text": ("Error: UI data for the current screen is not found in Data Engine."), + "image_path": None, + } + + try: + db_path = settings.DATA_ENGINE_DB_PATH + base_dir = settings.TRACES_PATH + storage = StorageManager(db_path=str(db_path), base_dir=base_dir) + + record = storage.get_ui_record(self.image_name) + if not record or not record.ocr_result: + return { + "text": ("No OCR results found in Data Engine for the current screen."), + "image_path": None, + } + + points = [] + labels = [] + text_output = [] + + for item in record.ocr_result: + text = item.get("text", "").strip() + bounds = item.get("bounds", []) + if text and len(bounds) == 4: + left, top, right, bottom = bounds + cx = (left + right) // 2 + cy = (top + bottom) // 2 + if 0 <= cx <= self.width and 0 <= cy <= self.height: + label_id = f"{prefix}{self.global_label_idx}" + self.global_label_idx += 1 + points.append([cx, cy]) + labels.append(label_id) + + cx_norm = int(max(0, min(1000, cx * 1000 / self.width))) + cy_norm = int(max(0, min(1000, cy * 1000 / self.height))) + + text_output.append(f"[{label_id}] '{text}' coords: [{cx_norm},{cy_norm}]") + + if not points: + return { + "text": "No valid OCR results found.", + "image_path": None, + } + + images_dir = base_dir / "images" + ocr_dir = images_dir / "ocr" + ocr_dir.mkdir(parents=True, exist_ok=True) + + output_path = ocr_dir / f"{self.image_name}_{self.global_label_idx}.jpg" + draw_dots(self.screenshot_path, points, labels, str(output_path), color=color) + + return { + "text": "\n".join(text_output), + "image_path": str(output_path), + } + + except Exception as e: + return { + "text": f"Error retrieving or processing OCR results: {e}", + "image_path": None, + } + + @trace(type="tool", name="inspect_region") + async def exec_inspect_region( + self, + x_min: int, + y_min: int, + x_max: int, + y_max: int, + zoom_factor: float = 2.0, + ) -> dict: + try: + px_start = int(max(0, min(self.width, x_min * self.width / 1000.0))) + py_start = int(max(0, min(self.height, y_min * self.height / 1000.0))) + px_end = int(max(0, min(self.width, x_max * self.width / 1000.0))) + py_end = int(max(0, min(self.height, y_max * self.height / 1000.0))) + + if px_end <= px_start or py_end <= py_start: + return { + "text": ( + f"Error: Invalid crop coordinates. x_min={x_min}," + f" x_max={x_max}, y_min={y_min}, y_max={y_max}." + " Bounding box dimensions must be strictly positive." + ), + "image_path": None, + } + + img = cv2.imread(self.screenshot_path) + if img is None: + return { + "text": ( + f"Error: Failed to read original screenshot at {self.screenshot_path}" + ), + "image_path": None, + } + + cropped = img[py_start:py_end, px_start:px_end] + + if zoom_factor is None: + zoom_factor = 2.0 + zoom_factor = max(1.0, min(4.0, float(zoom_factor))) + + c_h, c_w = cropped.shape[:2] + if c_h > 0 and c_w > 0: + new_w = int(round(c_w * zoom_factor)) + new_h = int(round(c_h * zoom_factor)) + if new_w > 0 and new_h > 0: + cropped = cv2.resize(cropped, (new_w, new_h)) + + base_dir = settings.TRACES_PATH + images_dir = base_dir / "images" + inspect_region_dir = images_dir / "inspect_region" + inspect_region_dir.mkdir(parents=True, exist_ok=True) + + image_name_safe = self.image_name or "temp_image" + output_path = inspect_region_dir / f"{image_name_safe}_{self.global_label_idx}.jpg" + cv2.imwrite(str(output_path), cropped) + + return { + "text": (f"Inspected region coordinates: [{x_min}, {y_min}] to [{x_max}, {y_max}]"), + "image_path": str(output_path), + } + + except Exception as e: + return { + "text": f"Error: Region inspection failed: {e}", + "image_path": None, + } + + async def _run_universal( + self, + query: str, + context_feedback: str, + screenshot_path: str, + state: State, + minimal_list: str, + version: str, + prompt_template: str, + max_iterations: int, + ) -> str: + """Executes Explorer reasoning loop via Universal LangChain ChatModel.""" + llm = get_llm(self.ctx, name="explorer") + + # Filter universal tools based on blacklist + exposed_tools = [ + t + for t in UNIVERSAL_EXPLORER_TOOLS + if t["function"]["name"] not in self.blacklisted_tools + ] + bound_llm = llm.bind_tools(exposed_tools) + + with open(screenshot_path, "rb") as f: + img_b64 = base64.b64encode(f.read()).decode("utf-8") + + user_content: list[dict[str, Any]] = [ + { + "type": "text", + "text": ( + f"Operator Request:\n- Query: {query}\n" + f"- Context Feedback: {context_feedback}\n\n" + "Initial marked UI elements list:\n" + f"{minimal_list}\n" + ), + }, + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{img_b64}"}, + }, + ] + + messages: list[BaseMessage] = [ + SystemMessage(content=prompt_template), + HumanMessage(content=user_content), + ] + + iterations = 0 + agent_outcome = "" + + while iterations < max_iterations: + iterations += 1 + is_final_turn = iterations == max_iterations + + if is_final_turn: + messages.append( + HumanMessage( + content=( + "[WARNING] This is your final iteration. You MUST" + " call 'submit_answer' to submit your final result." + ) + ) + ) + submit_only_tools = [ + t for t in exposed_tools if t["function"]["name"] == "submit_answer" + ] + current_llm = llm.bind_tools(submit_only_tools) if submit_only_tools else bound_llm + else: + current_llm = bound_llm + + response = await asyncio.wait_for(current_llm.ainvoke(messages), timeout=180) + messages.append(response) + + tool_calls = getattr(response, "tool_calls", None) or [] + if not tool_calls: + logger.warning( + f"Universal Explorer iteration {iterations}: No tool calls generated." + ) + if is_final_turn: + return json.dumps( + { + "candidates": [], + "fallback_message": str(response.content) or "No candidates found.", + }, + ensure_ascii=False, + ) + continue + + # Check for submit_answer + submit_call = next((tc for tc in tool_calls if tc.get("name") == "submit_answer"), None) + if submit_call: + args = submit_call.get("args", {}) + candidates = args.get("candidates", []) + fallback_message = args.get("fallback_message", "") + + # Validate coordinates + valid_candidates = [] + for cand in candidates: + if isinstance(cand, dict) and "coords" in cand: + coords = cand["coords"] + if isinstance(coords, list) and len(coords) == 2: + try: + nx, ny = int(coords[0]), int(coords[1]) + if 0 <= nx <= 1000 and 0 <= ny <= 1000: + valid_candidates.append(cand) + except (ValueError, TypeError): + pass + + agent_outcome = json.dumps( + { + "candidates": valid_candidates, + "fallback_message": fallback_message, + }, + ensure_ascii=False, + ) + return agent_outcome + + # Execute non-submit tools + for tc in tool_calls: + name = tc.get("name") + args = tc.get("args", {}) + call_id = tc.get("id", f"call_{iterations}_{name}") + + if name in self.blacklisted_tools: + messages.append( + ToolMessage( + tool_call_id=call_id, + name=name, + content=f"Tool '{name}' is blacklisted and unavailable.", + ) + ) + continue + + try: + if name == "ask_perception_tool": + res = await self.exec_ask_perception_tool( + search_query=args.get("search_query"), + nx=args.get("nx"), + ny=args.get("ny"), + detect_queries=args.get("detect_queries"), + ) + elif name == "detect_objects": + res = await self.exec_detect_objects( + queries=args.get("queries"), + target_image_id=args.get("target_image_id", "img_0"), + ) + elif name == "get_ocr_list": + res = await self.exec_get_ocr_list() + elif name == "ask_image_processor": + res = await self.exec_ask_image_processor( + instruction=args.get("instruction"), + target_image_id=args.get("target_image_id", "img_0"), + ) + elif name == "inspect_region": + res = await self.exec_inspect_region( + x_min=args.get("x_min"), + y_min=args.get("y_min"), + x_max=args.get("x_max"), + y_max=args.get("y_max"), + zoom_factor=args.get("zoom_factor", 2.0), + ) + else: + res = {"text": f"Error: Tool '{name}' is not recognized."} + + res_text = res.get("text") or res.get("result") or str(res) + messages.append(ToolMessage(tool_call_id=call_id, name=name, content=res_text)) + + except Exception as tool_err: + logger.error(f"Error executing tool '{name}': {tool_err}") + messages.append( + ToolMessage( + tool_call_id=call_id, + name=name, + content=f"Error executing tool '{name}': {tool_err}", + ) + ) + + if not agent_outcome: + agent_outcome = json.dumps( + { + "candidates": [], + "fallback_message": ( + "Explorer reached max iterations without finding candidates." + ), + }, + ensure_ascii=False, + ) + + return agent_outcome + + @trace(type="agent", name="explorer") + async def run( + self, + query: str, + context_feedback: str, + screenshot_path: str, + state: State, + minimal_list: str = "", + enable_caching: bool = False, + version: Literal["flash", "pro", "ultra"] = "pro", + ) -> str: + ctx = self.ctx + if version == "flash": + detector_prompt_path = ( + Path(__file__).parent.parent / "object_detector" / "object_detector.json" + ) + global_timeout = None + try: + with open(detector_prompt_path, encoding="utf-8") as f: + detector_config = json.load(f) + detector_templates = detector_config.get("templates", []) + detector_instructions = detector_config.get("instructions", "") + global_timeout = detector_config.get("global_timeout", None) + except Exception as e: + logger.warning(f"Failed to load detector prompt config: {e}") + detector_templates = [ + "Point to the following objects in the provided image: {labels_str}." + ] + detector_instructions = "" + + templates = [f"{t}\n\n{detector_instructions}" for t in 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 _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: + logger.error(f"Flash mode object detection failed: {e}") + return json.dumps( + { + "candidates": [], + "fallback_message": f"Flash mode detection error: {e}", + }, + ensure_ascii=False, + ) + + self.http_client = httpx.AsyncClient() + + # 1. Compute current image hash and verify with Data Engine + 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() + logger.info(f"Computed screenshot hash: {computed_hash}") + + # Check if it exists in DB + db_path = settings.DATA_ENGINE_DB_PATH + base_dir = settings.TRACES_PATH + storage = StorageManager(db_path, base_dir) + record = storage.get_image(computed_hash) + + if record: + image_name = computed_hash + else: + logger.warning( + f"Image hash {computed_hash} not found in Data Engine DB." + " Data Engine might not be synced yet." + ) + + except Exception as e: + logger.warning(f"Failed to compute hash or check DB: {e}") + + self.image_name = image_name + self.screenshot_path = screenshot_path + + # Resolve parameters + mode = version.capitalize() + version_blacklist = set() + if version == "ultra": + max_iterations = 8 + elif version == "pro": + max_iterations = 3 + version_blacklist = { + "ask_image_processor", + "get_ocr_list", + "inspect_region", + "detect_objects", + } + else: + max_iterations = 3 + version_blacklist = { + "ask_image_processor", + "get_ocr_list", + "inspect_region", + } + + try: + blacklisted_config = ( + self.ctx.agent_config.blacklisted_tools.get("explorer", []) + if self.ctx.agent_config + else [] + ) + self.blacklisted_tools = version_blacklist.union(set(blacklisted_config)) + except (TypeError, AttributeError): + self.blacklisted_tools = version_blacklist + + # 2. Prepare Native Tools declarations + 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 + + # Initialize Image Pool + 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 + + # 4. Initialize or reuse dynamic context-level GenAI client for connection pooling + client = getattr(ctx, "_genai_client", None) + if client is None: + logger.info( + "Initializing new GenAI client on context for connection pooling (Explorer)..." + ) + client = genai.Client( + api_key=settings.GOOGLE_API_KEY.get_secret_value() + if settings.GOOGLE_API_KEY + else None + ) + ctx._genai_client = client + + llm_cfg = 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.6-flash" + + # 5. Construct Prompt & Initial Message List + prompt_path = Path(__file__).parent / "explorer.json" + if not prompt_path.exists(): + return "Error: Explorer prompt template not found." + + try: + content = prompt_path.read_text(encoding="utf-8") + content_no_comments = re.sub(r"(? types.Part: + if use_file_api: + file_ref = client.files.upload(file=file_path) + uploaded_files.append(file_ref) + return types.Part( + file_data=types.FileData( + file_uri=file_ref.uri, + mime_type=file_ref.mime_type or "image/jpeg", + ) + ) + else: + with open(file_path, "rb") as f: + img_bytes = f.read() + return types.Part.from_bytes(data=img_bytes, mime_type="image/jpeg") + + cached_content = None + try: + # Upload initial screenshot + initial_part = get_image_part(img_to_read) + + contents = [ + types.Content( + role="user", + parts=[ + types.Part.from_text( + text=f"""Operator Request: +- Query: {query} +- Context Feedback: {context_feedback} + +Initial marked UI elements list (corresponding to numbers โ‘ , โ‘ก... in the image): +{minimal_list} +""" + ), + initial_part, + ], + ) + ] + + # Caching initialization + if enable_caching is None: + env_cache = os.getenv("ARTEMIS_EXPLORER_CACHING", "").lower() + if env_cache in ["true", "false"]: + enable_caching = env_cache == "true" + else: + enable_caching = getattr(settings, "EXPLORER_CACHING", True) + + cached_content = None + if enable_caching: + try: + logger.info("Creating cache resource for Explorer...") + cached_content = client.caches.create( + model=model_name, + config=types.CreateCachedContentConfig( + contents=[contents[0]], + system_instruction=prompt_template, + tools=[types.Tool(function_declarations=self.get_exposed_tools())], + tool_config=types.ToolConfig( + function_calling_config=types.FunctionCallingConfig( + mode=types.FunctionCallingConfigMode.ANY + ), + include_server_side_tool_invocations=True, + ), + ttl="300s", + ), + ) + logger.info(f"Cache resource created successfully: {cached_content.name}") + except Exception as cache_err: + logger.error(f"Failed to create cache resource: {cache_err}") + cached_content = None + + # 6. Execution Loop (Native SDK Tools Dispatching) + iterations = 0 + agent_outcome = "" + self.turn_latencies = [] + self.turn_cached_tokens = [] + self.trace_history = [] + + while iterations < max_iterations: + iterations += 1 + logger.info(f"Iteration {iterations}: Invoking Native Gemini SDK for Explorer...") + + self._prune_historical_images(contents, keep_last=1) + + is_final_turn = iterations == max_iterations + if is_final_turn: + warning_msg = ( + "\n[WARNING] This is your final iteration. You MUST" + " call 'submit_answer' to submit your final result. No" + " other tools are available." + ) + contents.append( + types.Content( + role="user", + parts=[types.Part.from_text(text=warning_msg)], + ) + ) + + start_turn = time.perf_counter() + with TraceSpan(name="gemini_explorer_call", ctx=ctx) as span: + if cached_content and not is_final_turn: + generate_config = types.GenerateContentConfig( + cached_content=cached_content.name, + temperature=temperature, + safety_settings=SAFETY_SETTINGS_BLOCK_NONE, + **( + { + "thinking_config": types.ThinkingConfig( + thinking_level=thinking_level + ) + } + if thinking_level + else {} + ), + ) + else: + generate_config = types.GenerateContentConfig( + system_instruction=prompt_template, + temperature=temperature, + safety_settings=SAFETY_SETTINGS_BLOCK_NONE, + tools=[ + types.Tool( + function_declarations=self.get_exposed_tools( + only_submit=is_final_turn + ) + ) + ], + tool_config=types.ToolConfig( + function_calling_config=types.FunctionCallingConfig( + mode=types.FunctionCallingConfigMode.ANY + ), + include_server_side_tool_invocations=True, + ), + **( + { + "thinking_config": types.ThinkingConfig( + thinking_level=thinking_level + ) + } + if thinking_level + else {} + ), + ) + + max_call_retries = 5 + for attempt in range(max_call_retries + 1): + try: + response = await asyncio.wait_for( + client.aio.models.generate_content( + model=model_name, + contents=contents, + config=generate_config, + ), + timeout=180, + ) + break + except APIError as api_err: + code = getattr(api_err, "code", None) + if code in [429, 503]: + if attempt >= max_call_retries: + raise api_err + backoff = (2.0**attempt) + random.uniform(0.1, 1.0) + logger.warning( + f"APIError {code} (overload/rate limit) on" + f" attempt {attempt + 1}. Retrying in" + f" {backoff:.2f}s..." + ) + await asyncio.sleep(backoff) + else: + raise api_err + except Exception as e: + if attempt >= max_call_retries: + raise e + backoff = (2.0**attempt) + random.uniform(0.1, 1.0) + logger.warning( + f"Unexpected error on attempt {attempt + 1}:" + f" {e}. Retrying in {backoff:.2f}s..." + ) + await asyncio.sleep(backoff) + end_turn = time.perf_counter() + turn_dur = end_turn - start_turn + self.turn_latencies.append(turn_dur) + + cached_tokens = 0 + if getattr(response, "usage_metadata", None) and getattr( + response.usage_metadata, + "cached_content_token_count", + None, + ): + cached_tokens = response.usage_metadata.cached_content_token_count + self.turn_cached_tokens.append(cached_tokens) + span.result = ( + f"Function calls: {len(response.function_calls)}" + if response.function_calls + else "Final answer" + ) + # Extract and log model thoughts and text if present + thinking_parts = [] + text_parts = [] + candidates = getattr(response, "candidates", None) + if not candidates and isinstance(response, dict): + candidates = response.get("candidates") + + if candidates and len(candidates) > 0: + candidate = candidates[0] + content = getattr(candidate, "content", None) + if not content and isinstance(candidate, dict): + content = candidate.get("content") + + parts = getattr(content, "parts", None) + if not parts and isinstance(content, dict): + parts = content.get("parts") + + if parts: + for part in parts: + is_thought = False + part_text = None + if isinstance(part, dict): + is_thought = part.get("thought", False) + part_text = part.get("text", None) + else: + is_thought = getattr(part, "thought", False) + part_text = getattr(part, "text", None) + + if is_thought and part_text: + thinking_parts.append(part_text) + elif part_text: + text_parts.append(part_text) + if thinking_parts: + span.payload["explorer_thought"] = "\n".join(thinking_parts) + if text_parts: + span.payload["explorer_text"] = "\n".join(text_parts) + + turn_record = { + "iteration": iterations, + "thoughts": ("\n".join(thinking_parts) if thinking_parts else ""), + "tool_calls": [], + } + + function_calls = response.function_calls + if not function_calls: + logger.warning( + f"Explorer iteration {iterations}: Model hallucinated" + " plain text. Forcing retry." + ) + turn_record["tool_calls"].append( + { + "name": "hallucinated_plain_text", + "args": {"text": "\n".join(text_parts) if text_parts else ""}, + "response": { + "error": ( + "Forced retry because model failed to call" + " submit_answer or any other tool" + ) + }, + } + ) + self.trace_history.append(turn_record) + contents.append(response.candidates[0].content) + contents.append( + types.Content( + role="user", + parts=[ + types.Part.from_text( + text=( + "You have not called 'submit_answer'" + " to submit your final answer yet. If" + " you are not certain about the final" + " answer, you can continue exploring" + f" {max_iterations - iterations} more" + " time(s)." + ) + ) + ], + ) + ) + continue + + # Record model's function call request in history + if response.candidates and response.candidates[0].content: + contents.append(response.candidates[0].content) + else: + contents.append( + types.Content( + role="model", + parts=[types.Part(function_call=fc) for fc in function_calls], + ) + ) + + # Execute function calls + tool_response_parts = [] + + # 6.1 Intercept and validate submit_answer to manage task lifecycle + submit_call = next( + ( + fc + for fc in function_calls + if (fc.name.split(":")[-1] if ":" in fc.name else fc.name) + == "submit_answer" + ), + None, + ) + if submit_call: + args = submit_call.args or {} + candidates = args.get("candidates", []) + errors = [] + + if len(function_calls) > 1: + errors.append( + "The tool 'submit_answer' MUST be called" + " individually in your final turn without any other" + " parallel tool calls. Please remove other tool" + " calls and re-submit." + ) + + fallback_message = args.get("fallback_message", "") + if not candidates and not fallback_message: + errors.append( + "The candidates list is empty. You must provide at" + " least one candidate matching the query, or" + " explain observations via fallback_message." + ) + + for i, cand in enumerate(candidates): + label = cand.get("label") + coords = cand.get("coords") + if not label: + errors.append(f"Candidate at index {i} is missing a 'label'.") + if not coords or not isinstance(coords, list) or len(coords) != 2: + errors.append( + f"Candidate '{label or i}' must have 'coords'" + " as an array of exactly 2 integers `[nx," + " ny]`." + ) + else: + try: + nx, ny = int(coords[0]), int(coords[1]) + if not (0 <= nx <= 1000) or not (0 <= ny <= 1000): + errors.append( + f"Candidate '{label or i}' coordinates" + f" `[{nx}, {ny}]` must strictly be in" + " the `[0-1000]` normalized scale" + " range inclusive." + ) + except (ValueError, TypeError, IndexError): + errors.append( + f"Candidate '{label or i}' coordinates are invalid integers." + ) + + if errors: + logger.warning(f"Explorer submit_answer validation failed: {errors}") + turn_record["tool_calls"].append( + { + "name": "submit_answer", + "args": args, + "response": {"error": ("Validation Failed:\n" + "\n".join(errors))}, + } + ) + tool_response_parts.append( + types.Part.from_function_response( + name="submit_answer", + response={ + "error": ( + "Validation Failed:\n" + + "\n".join(errors) + + "\nPlease correct these formatting" + " errors and re-submit using" + " submit_answer." + ) + }, + ) + ) + # Allow ReAct loop to continue so LLM can self-correct and re-submit + else: + logger.info("Explorer submit_answer validation passed successfully!") + agent_outcome = json.dumps(args, ensure_ascii=False) + turn_record["tool_calls"].append( + { + "name": "submit_answer", + "args": args, + "response": {"result": "success"}, + } + ) + self.trace_history.append(turn_record) + break + + for fc in function_calls: + name = fc.name.split(":")[-1] if ":" in fc.name else fc.name + if name == "submit_answer": + continue + args = fc.args or {} + tool_call_trace = { + "name": name, + "args": args, + } + if name in self.blacklisted_tools: + logger.warning( + "Explorer attempted to call blacklisted tool" + f" '{name}'. Blocking execution." + ) + tool_call_trace["response"] = { + "error": (f"Tool '{name}' is blacklisted and unavailable.") + } + turn_record["tool_calls"].append(tool_call_trace) + tool_response_parts.append( + types.Part.from_function_response( + name=fc.name, + response={ + "error": (f"Tool '{name}' is blacklisted and unavailable.") + }, + ) + ) + continue + + logger.info(f"Explorer executing tool '{name}' sequentially...") + + try: + if name == "ask_perception_tool": + res = await self.exec_ask_perception_tool( + search_query=args.get("search_query"), + nx=args.get("nx"), + ny=args.get("ny"), + detect_queries=args.get("detect_queries"), + ) + tool_call_trace["response"] = res + tool_response_parts.append( + types.Part.from_function_response( + name=name, + response={ + "text": res.get("text"), + }, + ) + ) + if res.get("image_paths"): + for img_p in res["image_paths"]: + try: + tool_response_parts.append(get_image_part(img_p)) + except Exception as e: + logger.warning( + f"Failed to load image response part for {img_p}: {e}" + ) + + elif name == "detect_objects": + res = await self.exec_detect_objects( + queries=args.get("queries"), + target_image_id=args.get("target_image_id", "img_0"), + ) + tool_call_trace["response"] = res + + if res.get("image_path"): + tool_response_parts.append( + types.Part.from_function_response( + name=name, + response={ + "result": res.get("text"), + }, + ) + ) + tool_response_parts.append(get_image_part(res["image_path"])) + else: + tool_response_parts.append( + types.Part.from_function_response( + name=name, + response={"result": res.get("text")}, + ) + ) + + elif name == "ask_image_processor": + res = await self.exec_ask_image_processor( + instruction=args.get("instruction"), + target_image_id=args.get("target_image_id", "img_0"), + ) + tool_call_trace["response"] = res + + tool_response_parts.append( + types.Part.from_function_response( + name=name, + response={"result": res.get("text")}, + ) + ) + if res.get("image_paths"): + for img_p in res["image_paths"]: + try: + tool_response_parts.append(get_image_part(img_p)) + except Exception as e: + logger.warning( + "Failed to upload image path" + f" {img_p} for ask_image_processor" + f" tool response: {e}" + ) + + elif name == "get_ocr_list": + res = await self.exec_get_ocr_list() + tool_call_trace["response"] = res + + if res.get("image_path"): + tool_response_parts.append( + types.Part.from_function_response( + name=name, + response={ + "result": res.get("text"), + }, + ) + ) + tool_response_parts.append(get_image_part(res["image_path"])) + else: + tool_response_parts.append( + types.Part.from_function_response( + name=name, + response={"result": res.get("text")}, + ) + ) + + elif name == "inspect_region": + res = await self.exec_inspect_region( + x_min=args.get("x_min"), + y_min=args.get("y_min"), + x_max=args.get("x_max"), + y_max=args.get("y_max"), + zoom_factor=args.get("zoom_factor"), + ) + tool_call_trace["response"] = res + + if res.get("image_path"): + tool_response_parts.append( + types.Part.from_function_response( + name=name, + response={ + "result": res.get("text"), + }, + ) + ) + tool_response_parts.append(get_image_part(res["image_path"])) + else: + tool_response_parts.append( + types.Part.from_function_response( + name=name, + response={"result": res.get("text")}, + ) + ) + else: + tool_call_trace["response"] = {"error": f"Tool {name} not found"} + tool_response_parts.append( + types.Part.from_function_response( + name=name, + response={"error": f"Tool {name} not found"}, + ) + ) + except Exception as e: + logger.error(f"Explorer tool {name} execution failed: {e}") + tool_call_trace["response"] = {"error": str(e)} + tool_response_parts.append( + types.Part.from_function_response(name=name, response={"error": str(e)}) + ) + finally: + turn_record["tool_calls"].append(tool_call_trace) + + if tool_response_parts: + # Native Gemini SDK allows 'tool' role to return mixed parts + contents.append(types.Content(role="user", parts=tool_response_parts)) + + self.trace_history.append(turn_record) + + if iterations >= max_iterations and not agent_outcome: + agent_outcome = ( + "Error: Explorer reached maximum iterations without a conclusive answer." + ) + + except Exception as e: + logger.error(f"Explorer execution loop failed: {e}") + + err_msg = str(e) + if "preempted" in err_msg.lower(): + clean_msg = ( + "Model request was preempted by the server due to high" + " demand. Please try again later." + ) + elif "overloaded" in err_msg.lower() or "503" in err_msg: + clean_msg = "Model service is currently overloaded. Please try again later." + elif "quota" in err_msg.lower() or "429" in err_msg: + clean_msg = "Model API quota exceeded or rate limited." + elif "thinkingconfig" in err_msg.lower(): + clean_msg = ( + "Model configuration error: ThinkingConfig is not" + f" supported for model {model_name}." + ) + else: + clean_msg = f"Explorer agent execution failed: {err_msg}" + + agent_outcome = json.dumps( + {"candidates": [], "fallback_message": clean_msg}, + ensure_ascii=False, + ) + finally: + if self.http_client: + try: + await self.http_client.aclose() + logger.info("Closed Explorer HTTP client.") + except Exception as close_err: + logger.warning(f"Failed to close Explorer HTTP client: {close_err}") + if cached_content: + try: + logger.info(f"Deleting cache resource: {cached_content.name}") + client.caches.delete(name=cached_content.name) + except Exception as cleanup_cache_err: + logger.warning(f"Failed to delete cache resource: {cleanup_cache_err}") + for file_ref in uploaded_files: + try: + client.files.delete(name=file_ref.name) + except Exception as cleanup_err: + logger.warning(f"Failed to delete uploaded file {file_ref.name}: {cleanup_err}") + logger.info(f"Explorer turn latencies: {self.turn_latencies}") + + return agent_outcome diff --git a/artemis/agents/flash/context_compressor.py b/artemis/agents/flash/context_compressor.py new file mode 100644 index 0000000..a8805e9 --- /dev/null +++ b/artemis/agents/flash/context_compressor.py @@ -0,0 +1,118 @@ +# 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. + +"""Universal Context Compressor for Artemis Flash profile. + +Replaces pruned historical screenshots with high-density visual step summaries +when ready, while seamlessly falling back to standard pruned action text if +the background summary is still in flight. +""" + +from typing import Any + +from langchain_core.messages import BaseMessage + +from artemis.agents.flash.summarizer import VisualStepSummarizer +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +def compress_flash_messages( + messages: list[BaseMessage], + summarizer: VisualStepSummarizer | None = None, + prune_history_xml: bool = True, +) -> None: + """Compresses multi-turn conversation messages in-place. + + Rules: + 1. Keeps the latest screenshot and live UI hierarchy completely intact. + 2. For earlier steps, if a VisualStepSummarizer has completed the step's summary, + replaces the intermediate image block with the objective summary text. + 3. If the background summary is not ready yet, gracefully falls back to the existing + behavior (omitting the image block and retaining the original action text). + 4. If prune_history_xml is True, removes heavy outdated UI Element lists from past turns. + """ + if not messages: + return + + # 1. Identify the last message containing an active screenshot + last_img_msg_idx = -1 + for idx in range(len(messages) - 1, -1, -1): + msg = messages[idx] + msg_content = getattr(msg, "content", None) + if isinstance(msg_content, list): + has_image = any( + isinstance(block, dict) and block.get("type") in ("image_url", "image") + for block in msg_content + ) + if has_image: + last_img_msg_idx = idx + break + + # 2. Iterate through historical messages and apply compression/fallback + tool_step_counter = 0 + for idx in range(len(messages)): + msg = messages[idx] + msg_content = getattr(msg, "content", None) + if not isinstance(msg_content, list): + continue + + # Keep the latest live screen intact + if idx == last_img_msg_idx: + continue + + # Count tool steps strictly corresponding to FlashRunner action turns + is_tool_msg = getattr(msg, "tool_call_id", None) is not None + if is_tool_msg: + tool_step_counter += 1 + current_step_num = tool_step_counter + else: + current_step_num = None + + new_blocks: list[dict[str, Any]] = [] + summary_injected = False + + for b in msg_content: + if not isinstance(b, dict): + continue + b_type = b.get("type", "") + b_text = b.get("text", "") + + # Filter heavy outdated XML hierarchy from past steps to save tokens + if prune_history_xml and b_type == "text" and "--- UI Element List ---" in b_text: + continue + + if b_type in ("image_url", "image"): + # Check if this is a tool execution step and background summarizer has a ready summary + if ( + current_step_num + and summarizer + and summarizer.has_summary(current_step_num) + and not summary_injected + ): + summary_text = summarizer.get_summary(current_step_num) + if summary_text: + new_blocks.append({"type": "text", "text": summary_text}) + summary_injected = True + # If summarizer is not ready or initial screen, gracefully omit the image block + continue + + new_blocks.append(b) + + # Defensive fallback: never leave ToolMessage or HumanMessage content empty + if not new_blocks: + new_blocks = [{"type": "text", "text": "Action completed." if is_tool_msg else ""}] + + msg.content = new_blocks diff --git a/artemis/agents/flash/flash_runner.md b/artemis/agents/flash/flash_runner.md new file mode 100644 index 0000000..0618206 --- /dev/null +++ b/artemis/agents/flash/flash_runner.md @@ -0,0 +1,41 @@ +# ROLE & OBJECTIVE +You are **ARTEMIS Flash Runner**, an autonomous and highly efficient Android Device Execution Agent. Your goal is to accomplish the user's objective on the device. + +**Objective: {{ goal }}** + +--- + +# 1. COGNITIVE PROTOCOL +You operate in a dynamic observation-reasoning-execution loop. **CRITICAL RULE: YOU MUST NEVER INVOKE A TOOL SILENTLY!** In every single turn, you **MUST FIRST output a natural language reasoning paragraph BEFORE generating any tool call.** + +- **Turn 1 Only (Initial Planning)**: Briefly decompose the overall `{{ goal }}` into high-level visual **Milestones** in your text output. Do not repeat this macro breakdown in subsequent turns. + +--- + +# 2. DYNAMIC ENVIRONMENT & TIME AWARENESS +Keep in mind that the mobile device is an asynchronous, constantly evolving environment. However, your actions entail an inherent "turn latency" (a 3โ€“7 second delay caused by model inference and network transmission between normal consecutive actions). +- **State Drift & Progression Tolerance**: While you think, video playback advances, banners scroll, and web pages finish loading. If the latest observation has progressed further than expected (e.g., an ad automatically skipped or a redirect landed), **do not treat it as an error**. Adapt smoothly to the live screen in front of you. +- **Timed Waiting & Transitions**: When facing mid-transition animations, loading spinners, required watch durations, or ad countdown timers, use the precise `wait_for_delay` command to allow time to pass or the UI state to stabilize. Some transition states are subtle and may even resemble a phone bug or system freeze; it is always advisable to wait first, and only consider retrying or choosing an alternative path if the state remains broken. +- **Time-Sensitive Tasks & Transient UI**: First, you must understand that each of your decision turns takes several seconds of model inference and network delay, while the mobile device state is continuously changing. A target element that you see on the current screenshot (e.g., auto-dismissing overlays, temporary message bars, menus or control bars with short timeouts) may have already vanished by the time your physical click command reaches the device! This is one of the most common reasons why the UI does not change as expected after a single step. We call these **Time-Sensitive Tasks**. + - *Mental Model*: When facing these tasks, any attempt to interact step-by-step (such as tapping once to wake up a menu/bar and waiting for the next turn to tap the target button) WILL fail because the element will auto-dismiss during your decision delay. You MUST combine the summon action on the trigger element (`Trigger`) and the target action (`Target`) into a single atomic sequence. + - *Chained Execution*: Use `click_sequence(sequence=[[trigger_x, trigger_y], [target_x, target_y]])` in `0-1000` normalized scale to execute both taps consecutively with a 50ms delay, catching the transient element before the timeout closes it. + +--- + +# 3. STAGNATION & LOOP PREVENTION +- **Retrying is not the only option.**: If you perform the same action on the same target for two consecutive turns and the screen or XML shows **no substantial change**, further retries are unlikely to be useful; the probability of the execution tool itself malfunctioning is very low, so you should pause and consider the underlying reason. +- **Flexibly Adjust Strategies:** Always prioritize resolving the immediate issue. If you encounter overlays or pop-ups blocking input, or if an element is unresponsive, actively consider alternative ways to complete the task. If an action fails to proceed as expected, systematically consider whether the issue lies with the device itself, the operational path, or a change in state occurring between your actual interaction and the screenshot (Time-Sensitive Tasks). + +--- + +# 4. ACTION & MEMORY RULES +1. **Normalized Coordinates (`0-1000` Scale)**: Always use normalized `[x, y]` coordinates in `0-1000` scale for any location-based target (e.g., `target=[500, 600]` or `click_sequence=[[300, 400], [500, 600]]`). +2. **Do not determine element coordinates based on guesswork.** If the element is visible in the screenshot but its coordinates are missing from the provided XML tree, please use the `ask_explorer` tool to locate the element first. +3. **Screenshot Pruning Awareness**: To optimize speed, the system **prunes old intermediate screenshots from your history**โ€”you only ever see the *latest* live screenshot and UI tree. Therefore, if you spot critical text in an earlier step that you will need later (e.g., verification codes, specific titles, search lists), **you must explicitly verbalize and write those details into your text thinking/notes** so they persist across turns. +4. **App Launching**: To start or reset an application, always use `manage_app(action="launch", app_name="...")` or `action="stop"`. Do not manually swipe across home screens. + +--- + +# 5. TERMINATION PROTOCOL (`report_task_status`) +- **Completion (`status="completed"`)**: ONLY call `report_task_status(status="completed", ...)` after you have **visually verified** from the latest screen that the user's objective (`{{ goal }}`) is fully achieved. +- **Failure (`status="failed"`)**: If the task is genuinely impossible due to paywalls, unrecoverable crashes, or hitting an absolute dead end after sensible fallbacks, call `report_task_status(status="failed", ...)` with the exact blocking reason. diff --git a/artemis/agents/flash/flash_summarizer.md b/artemis/agents/flash/flash_summarizer.md new file mode 100644 index 0000000..f00587d --- /dev/null +++ b/artemis/agents/flash/flash_summarizer.md @@ -0,0 +1,41 @@ +# ROLE & OBJECTIVE +You are the Step Summarizer for an Android UI automation agent. Your task is to synthesize the Before/After screenshots and physical action into a concise, high-information-density, strictly objective first-person historical memory. + +--- + +# PERSPECTIVE & FORMAT CONSTRAINTS +1. **First-Person Perspective**: Write strictly from the agent's first-person perspective using **"I"** (e.g., "In Step {{ step_number }}, I tapped...", "I swiped up on... and observed..."). NEVER use third-person terms like "The agent", "The operator", or "The system". +2. **Single Continuous Paragraph**: Your output MUST be exactly **one compact, continuous paragraph** (1โ€“3 sentences, 35โ€“65 words). +3. **No Lists or Formatting**: Do NOT use bullet points, numbered lists, markdown headers, bold labels, or line breaks in your output. + +--- + +# CORE OBSERVATION RULES +1. **Zero Subjective Validation (STRICT)**: + - Absolutely avoid declaring semantic goal achievement, subjective success, or final failure. In historical compression, subjective assumptions cause future turns to falsely believe a subgoal was definitively accomplished or permanently blocked. + - **STRICTLY BANNED WORDS**: `successfully`, `completed`, `achieved`, `entered`, `navigated to`, `arrived at`, `failed`, `unsuccessful`, `could not`, `impossible`. +2. **Action & Visual Delta Structure**: + - **Target & Action**: Identify what control/button was targeted (referencing the red visual indicator on the BEFORE screen) and the physical action performed. + - **Objective Transition**: Describe the physical screen change (e.g., page transitioned, modal dialog opened, list scrolled revealing new items, checkbox/toggle toggled, keyboard appeared). + - **Preserve Critical Data & Verifications**: Faithfully transcribe visible text, toast alerts, error banners, verification codes, prices, account names, or tracking numbers that appeared on the AFTER screen. + +--- + +# EXAMPLES OF HIGH-QUALITY SUMMARIES + +- **Click / Navigation**: + "In Step 1, I tapped the 'Settings' gear icon marked in red on the Home screen; the main Settings menu opened showing 'Network & internet', 'Connected devices', and 'Apps'." + +- **Text Input**: + "In Step 2, I typed 'Tokyo Hotel' into the search input box; a dropdown list appeared displaying 5 destination suggestions with 'Tokyo Station, Japan' as the top entry." + +- **Swipe / Scroll**: + "In Step 3, I swiped up along the flight results list; the screen scrolled downward approximately one page, revealing 3 additional evening flights starting from '$420'." + +- **Modal / Alert State**: + "In Step 4, I tapped the 'Confirm Booking' button; an alert dialog titled 'Payment Method Required' appeared over the view with an 'Add Card' option." + +--- + +# TASK +Generate the single-paragraph first-person summary for Step {{ step_number }} following all constraints above: diff --git a/artemis/agents/flash/runner.py b/artemis/agents/flash/runner.py new file mode 100644 index 0000000..939d0d6 --- /dev/null +++ b/artemis/agents/flash/runner.py @@ -0,0 +1,524 @@ +# 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. + +"""Universal Multi-Model FlashRunner for Artemis. + +Executes autonomous reactive mobile workflows across Google Gemini, +OpenAI GPT-4o/o3, Anthropic Claude 3.5/3.7, and OpenRouter endpoints. +""" + +import asyncio +import base64 +import json +from pathlib import Path +import uuid + +from jinja2 import Template +from langchain_core.messages import ( + BaseMessage, + HumanMessage, + SystemMessage, + ToolMessage, +) + +from artemis.agents.flash.context_compressor import compress_flash_messages +from artemis.agents.flash.summarizer import VisualStepSummarizer +from artemis.agents.validator.tool_declarations import ( + ACTION_TOOL_NAMES, + ASK_EXPLORER_TOOL, + CLICK_SEQUENCE_TOOL, + MobileActionExecutor, + REPORT_TASK_STATUS_TOOL, + VALIDATOR_TOOLS_DECLARATION, + capture_screenshot_and_parse_ui, + prune_intermediate_screenshots, +) +from artemis.config import StepSummarizerConfig, load_agent_config +from artemis.context import ArtemisContext +from artemis.controllers.unified_controller import UnifiedMobileController +from artemis.data_engine.trace import trace +from artemis.graph.perception import _check_injected_instruction_file +from artemis.graph.state import State +from artemis.services.llm import ( + RobustChatModelWrapper, + get_google_llm, + get_llm, + invoke_llm_with_timeout_message, +) +from artemis.utils.coordinates import parse_swipe_parameters +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +class FlashRunner: + """Reactive, ultra-fast agent execution loop supporting all VLM providers.""" + + def __init__(self, ctx: ArtemisContext, goal: str, max_turns: int | None = None): + self.ctx = ctx + self.goal = goal + try: + cfg = load_agent_config() + self.max_turns = max_turns if max_turns is not None else cfg.flash.max_turns + self.step_summarizer_cfg = cfg.flash.step_summarizer + except Exception: + self.max_turns = max_turns if max_turns is not None else 30 + self.step_summarizer_cfg = StepSummarizerConfig() + + self.controller = UnifiedMobileController(ctx) + self.executor = MobileActionExecutor(ctx, self.controller) + self.summarizer = ( + VisualStepSummarizer(ctx, model_name=self.step_summarizer_cfg.model) + if self.step_summarizer_cfg.enabled + else None + ) + + def _get_tools(self) -> list: + tools = [t for t in VALIDATOR_TOOLS_DECLARATION if t.name != "report_failure_analysis"] + tools.insert(1, CLICK_SEQUENCE_TOOL) + tools.append(ASK_EXPLORER_TOOL) + tools.append(REPORT_TASK_STATUS_TOOL) + return tools + + def _prune_intermediate_screenshots(self, messages: list[BaseMessage]) -> None: + """Prunes binary screenshot blocks from intermediate observation messages.""" + prune_intermediate_screenshots(messages) + + @trace(type="agent", name="FlashRunner") + async def run(self, state: State) -> dict: + logger.info(f"Starting Artemis Flash reactive loop for goal: {self.goal}") + + # 1. Initialize Universal LLM via Service Layer + try: + llm = get_llm(self.ctx, name="operator") + except Exception as e: + logger.warning(f"Failed to get operator LLM from config, using default: {e}") + + llm = RobustChatModelWrapper(get_google_llm(model_name="gemini-2.5-flash"), self.ctx) + + tools_declaration = self._get_tools() + + # 2. Capture Initial State (Screenshot + UI Tree) + shot_path, img_bytes, xml_list = await capture_screenshot_and_parse_ui( + self.ctx, state, self.controller, skip_settling=False + ) + state.latest_screenshot = shot_path + + user_content: list[dict] = [ + {"type": "text", "text": f"Your objective is: {self.goal}"}, + ] + if img_bytes: + img_b64 = base64.b64encode(img_bytes).decode("utf-8") + user_content.append( + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{img_b64}"}, + } + ) + if xml_list: + user_content.append( + { + "type": "text", + "text": f"--- UI Element List ---\n{xml_list}", + } + ) + user_content.append( + { + "type": "text", + "text": ( + "CRITICAL RULE: In every single turn, you MUST FIRST output a natural" + " language reasoning/explanation paragraph BEFORE invoking any tool call." + ), + } + ) + + # Render System Prompt + prompt_path = Path(__file__).parent / "flash_runner.md" + prompt_template = prompt_path.read_text(encoding="utf-8") + system_prompt = Template(prompt_template).render(goal=self.goal) + + messages: list[BaseMessage] = [ + SystemMessage(content=system_prompt), + HumanMessage(content=user_content), + ] + + turns = 0 + final_report = None + current_pre_screenshot_bytes = img_bytes + current_xml_list = xml_list + + while turns < self.max_turns: + turns += 1 + logger.info(f"--- Artemis Flash Turn {turns}/{self.max_turns} ---") + + if self.ctx.data_engine: + self.ctx.data_engine.allocate_step_id() + + # Check for real-time injected instructions + if self.ctx.data_engine and self.ctx.data_engine.base_dir: + try: + injected_instruction = await asyncio.to_thread( + _check_injected_instruction_file, + str(self.ctx.data_engine.base_dir), + ) + if injected_instruction: + messages.append( + HumanMessage( + content=( + "[REAL-TIME INJECTED INSTRUCTION from user]:" + f" {injected_instruction}\nYou MUST immediately follow" + " this instruction and adjust your plan/actions." + ) + ) + ) + except Exception as e: + logger.warning(f"Failed to check injected instruction in FlashRunner: {e}") + + # Compress previous history before LLM call to save context and inject objective summaries + if self.summarizer: + compress_flash_messages( + messages, + summarizer=self.summarizer, + prune_history_xml=self.step_summarizer_cfg.prune_history_xml, + ) + else: + prune_intermediate_screenshots(messages) + + # Tool restriction on the final turn + if turns == self.max_turns: + messages.append( + HumanMessage( + content=( + "[WARNING] This is your final turn; only" + " 'report_task_status' is available." + ) + ) + ) + current_tools = [t for t in tools_declaration if t.name == "report_task_status"] + else: + current_tools = tools_declaration + + # Bind active tools + bound_llm = llm.bind_tools(current_tools) + + # Invoke Model with Streaming & Retry + response = None + max_retries = 3 + + for attempt in range(max_retries): + try: + + async def run_stream(): + full_response = None + async for chunk in bound_llm.astream(messages): + if full_response is None: + full_response = chunk + else: + full_response += chunk + return full_response + + response = await invoke_llm_with_timeout_message( + run_stream(), timeout_seconds=10, hard_timeout=180 + ) + if response: + break + except Exception as e: + logger.warning(f"FlashRunner LLM turn failed on attempt {attempt + 1}: {e}") + if attempt == max_retries - 1: + raise e + await asyncio.sleep(2.0**attempt) + + if response is None: + break + + messages.append(response) + + # Extract thought and tool calls + raw_text = response.content if isinstance(response.content, str) else "" + if isinstance(response.content, list): + raw_text = "".join( + b.get("text", "") + for b in response.content + if isinstance(b, dict) and "text" in b + ) + + # Extract token usage metadata from response + step_token_usage = None + if hasattr(response, "usage_metadata") and response.usage_metadata: + u = response.usage_metadata + pr = u.get("input_tokens") or u.get("prompt_tokens") or 0 + co = u.get("output_tokens") or u.get("completion_tokens") or 0 + to = u.get("total_tokens") or (pr + co) + if to > 0: + step_token_usage = { + "prompt_tokens": int(pr), + "completion_tokens": int(co), + "total_tokens": int(to), + } + elif hasattr(response, "response_metadata") and isinstance( + response.response_metadata, dict + ): + u = response.response_metadata.get( + "usage_metadata" + ) or response.response_metadata.get("token_usage") + if isinstance(u, dict): + pr = ( + u.get("input_tokens") + or u.get("prompt_tokens") + or u.get("prompt_token_count") + or 0 + ) + co = ( + u.get("output_tokens") + or u.get("completion_tokens") + or u.get("candidates_token_count") + or 0 + ) + to = u.get("total_tokens") or u.get("total_token_count") or (pr + co) + if to > 0: + step_token_usage = { + "prompt_tokens": int(pr), + "completion_tokens": int(co), + "total_tokens": int(to), + } + + if not step_token_usage or step_token_usage.get("total_tokens", 0) <= 0: + calc_prompt_chars = 0 + calc_images = 0 + for msg in messages: + c = getattr(msg, "content", "") + if isinstance(c, str): + calc_prompt_chars += len(c) + elif isinstance(c, list): + for block in c: + if isinstance(block, dict): + if block.get("type") == "image_url" or "image_url" in block: + calc_images += 1 + else: + calc_prompt_chars += len(str(block.get("text", ""))) + else: + calc_prompt_chars += len(str(block)) + prompt_tokens = (calc_prompt_chars // 4) + (calc_images * 258) + completion_tokens = len(raw_text) // 4 + step_token_usage = { + "prompt_tokens": max(1, prompt_tokens), + "completion_tokens": max(1, completion_tokens), + "total_tokens": max(1, prompt_tokens + completion_tokens), + } + + if self.ctx.data_engine: + current_step_id = getattr(self.ctx.data_engine, "current_step_id", None) + self.ctx.data_engine.record_trace( + type="llm_call", + name="FlashRunner", + payload={"token_usage": step_token_usage, "response": raw_text}, + step_id=current_step_id, + status="success", + ) + + tool_calls = response.tool_calls or [] + + # Fallback text parsing if tool_calls not parsed natively + if not tool_calls and "```json" in raw_text: + try: + snippet = raw_text.split("```json")[1].split("```")[0].strip() + parsed = json.loads(snippet) + if isinstance(parsed, dict) and "name" in parsed: + tool_calls = [ + { + "name": parsed["name"], + "args": parsed.get("args", {}), + "id": str(uuid.uuid4()), + } + ] + except Exception: + pass + + if not tool_calls: + logger.info( + f"FlashRunner received response without tool calls at turn {turns}:" + f" {raw_text[:100]}..." + ) + if turns == self.max_turns: + final_report = {"status": "failed", "explanation": raw_text} + break + messages.append( + HumanMessage( + content=( + "You did not call any tools. Please make progress" + " by calling an action tool or 'report_task_status'." + ) + ) + ) + continue + + # Process tool calls + for tc in tool_calls: + name = tc["name"].split(":")[-1] if ":" in tc["name"] else tc["name"] + args = tc.get("args") or {} + tc_id = tc.get("id") or str(uuid.uuid4()) + logger.info(f"Executing Flash tool: {name}({args})") + + if name == "report_task_status": + final_report = args + if self.ctx.data_engine: + try: + if self.ctx.data_engine.current_step_id is None: + self.ctx.data_engine.allocate_step_id() + self.ctx.data_engine.record_step( + pre_screenshot_bytes=current_pre_screenshot_bytes, + ui_tree=current_xml_list, + action_taken={"action": "report_task_status", "args": args}, + operator_raw_thinking=raw_text, + last_execution_result={ + "result": "Task completed with final report." + }, + extra_metadata={"token_usage": step_token_usage}, + ) + except Exception as step_err: + logger.warning(f"Error recording final step in FlashRunner: {step_err}") + messages.append( + ToolMessage( + tool_call_id=tc_id, + name=name, + content=json.dumps({"status": "acknowledged"}), + status="success", + ) + ) + if self.summarizer: + await self.summarizer.flush() + return final_report + + try: + exec_result = await self.executor.execute(name, args, tc_id, state) + + post_img_bytes = exec_result.screenshot_bytes + if not post_img_bytes and name in ACTION_TOOL_NAMES: + try: + controller = UnifiedMobileController(self.ctx) + screen_data = await controller.get_screen_data() + post_img_bytes = base64.b64decode(screen_data.base64) + if not exec_result.ui_elements_text: + exec_result.ui_elements_text = screen_data.elements + except Exception as shot_err: + logger.warning( + f"Failed to capture fallback screenshot in FlashRunner: {shot_err}" + ) + + # Record telemetry / step in DataEngine + if self.ctx.data_engine and name in ACTION_TOOL_NAMES: + try: + if self.ctx.data_engine.current_step_id is None: + self.ctx.data_engine.allocate_step_id() + + # Extract and enrich coordinate metadata + norm_coords = None + norm_start = None + norm_end = None + if name == "swipe": + kind, target_val, _ = parse_swipe_parameters(args) + if ( + kind == "coords" + and isinstance(target_val, list) + and len(target_val) == 4 + ): + norm_coords = target_val + norm_start = target_val[:2] + norm_end = target_val[2:] + elif kind == "direction" and isinstance(target_val, str): + g_lower = target_val.lower() + if "up" in g_lower: + norm_coords = [600, 800, 600, 200] + elif "down" in g_lower: + norm_coords = [600, 200, 600, 800] + elif "left" in g_lower: + norm_coords = [800, 500, 200, 500] + elif "right" in g_lower: + norm_coords = [200, 500, 800, 500] + if norm_coords: + norm_start = norm_coords[:2] + norm_end = norm_coords[2:] + elif name in ("click", "tap", "long_press", "input_text"): + target = args.get("target") or args.get("coordinates") + if isinstance(target, (list, tuple)) and len(target) == 2: + norm_coords = list(target) + elif isinstance(target, str): + nums = re.findall(r"-?\d+(?:\.\d+)?", target) + if len(nums) == 2: + norm_coords = [int(float(nums[0])), int(float(nums[1]))] + + action_dict = { + "action": name, + "coordinates": ( + args.get("target") + or args.get("coordinates") + or args.get("sequence") + or norm_coords + ), + "args": args, + } + if norm_coords: + action_dict["normalized_coordinates"] = norm_coords + if norm_start and norm_end: + action_dict["normalized_start_coordinates"] = norm_start + action_dict["normalized_end_coordinates"] = norm_end + + self.ctx.data_engine.record_step( + pre_screenshot_bytes=current_pre_screenshot_bytes, + post_screenshot_bytes=post_img_bytes, + ui_tree=(exec_result.ui_elements_text or current_xml_list), + action_taken=action_dict, + operator_raw_thinking=raw_text, + last_execution_result={"result": exec_result.text_summary}, + extra_metadata={"token_usage": step_token_usage}, + ) + except Exception as step_err: + logger.warning(f"Error recording step in FlashRunner: {step_err}") + + # โšก Non-blocking dispatch of objective visual transition summarizer + if self.summarizer and name in ACTION_TOOL_NAMES: + self.summarizer.dispatch( + step_number=turns, + action_name=name, + action_args=args, + pre_img_bytes=current_pre_screenshot_bytes, + post_img_bytes=post_img_bytes, + exec_outcome=exec_result.text_summary, + ) + + if post_img_bytes: + current_pre_screenshot_bytes = post_img_bytes + if exec_result.ui_elements_text: + current_xml_list = exec_result.ui_elements_text + + messages.append(exec_result.to_langchain_tool_message()) + + except Exception as e: + logger.error(f"Error executing tool {name}: {e}") + messages.append( + ToolMessage( + tool_call_id=tc_id, + name=name, + content=f"Error executing tool {name}: {e}", + status="error", + ) + ) + + if self.summarizer: + await self.summarizer.flush() + + return final_report or { + "status": "failed", + "explanation": "Max turns reached without final status report.", + } diff --git a/artemis/agents/flash/summarizer.py b/artemis/agents/flash/summarizer.py new file mode 100644 index 0000000..20a0095 --- /dev/null +++ b/artemis/agents/flash/summarizer.py @@ -0,0 +1,226 @@ +# 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) diff --git a/artemis/agents/history_analyzer/history_analyzer.md b/artemis/agents/history_analyzer/history_analyzer.md new file mode 100644 index 0000000..c959941 --- /dev/null +++ b/artemis/agents/history_analyzer/history_analyzer.md @@ -0,0 +1,8 @@ +# IDENTITY +You are a **History Analyzer**, expert in analyzing multi-agent session execution history to answer user queries. + +# MANDATE +1. You are provided with the **Task Plan and Execution History** summarizing the high-level plan and history steps. +2. For queries about specific step details (such as precise actions, coordinates, `operator_raw_thinking`, or `last_execution_result`), you should call the `get_step_details` tool. +3. Never assume, guess, or invent details not visible in the Task Plan and Execution History. Always query `get_step_details` for specifics. If details are unavailable, state it clearly. +4. Make good use of note tools to see what notes are available. When analyzing history records, open and cross-reference relevant notes to ensure you have grasped all necessary information. diff --git a/artemis/agents/history_analyzer/history_analyzer.py b/artemis/agents/history_analyzer/history_analyzer.py new file mode 100644 index 0000000..3689951 --- /dev/null +++ b/artemis/agents/history_analyzer/history_analyzer.py @@ -0,0 +1,250 @@ +# 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 json +from pathlib import Path +from typing import Any + +from langchain_core.messages import ( + BaseMessage, + HumanMessage, + SystemMessage, + ToolMessage, +) +from langchain_core.tools import BaseTool, tool +from artemis.context import ArtemisContext +from artemis.data_engine.trace import CURRENT_TRACE_ID, TraceSpan, trace +from artemis.services.llm import get_llm, invoke_llm_with_timeout_message +from artemis.tools.scratchpad import ( + get_list_notes_tool_pure, + get_read_note_tool_pure, +) +from artemis.utils.logger import get_logger +from artemis.utils.notes import ( + get_note_file_path, + get_notes_dir, +) +from artemis.utils.task_tree import build_plan_and_history, get_active_subgoal_hashes + +logger = get_logger(__name__) + + +class HistoryAnalyzer: + def __init__(self, ctx: ArtemisContext): + self.ctx = ctx + + def exec_get_step_details(self, start_step: int, end_step: int) -> str: + """Retrieve the detailed, full-granularity information for a range of steps (inclusive). + + Use this when you need to inspect specific actions taken, operator + thinking, or execution results. + """ + try: + s_step = int(start_step) + e_step = int(end_step) + except (ValueError, TypeError): + return ( + f"Error: start_step and end_step must be integers, got {start_step} and {end_step}." + ) + + matched_steps = [] + for s in getattr(self, "history_steps", []): + step_num = s.get("step_number") + if step_num is not None and s_step <= step_num <= e_step: + details = { + "step_number": s.get("step_number"), + "relative_time": s.get("relative_time"), + "summary": s.get("summary"), + "action_taken": s.get("action_taken"), + "operator_raw_thinking": s.get("operator_raw_thinking"), + "last_execution_result": s.get("last_execution_result"), + } + matched_steps.append(details) + if not matched_steps: + return f"No steps found in range [{s_step}, {e_step}]." + return json.dumps(matched_steps, indent=2, ensure_ascii=False) + + def _get_step_details_tool(self, history_steps: list[dict[str, Any]]) -> BaseTool: + self.history_steps = history_steps + + @tool + def get_step_details(start_step: int, end_step: int) -> str: + """Retrieve the detailed, full-granularity information for a range of steps (inclusive). + + Use this when you need to inspect specific actions taken, operator + thinking, or execution results. + """ + return self.exec_get_step_details(start_step, end_step) + + return get_step_details + + @trace(type="agent", name="history_analyzer") + async def run(self, query: str) -> str: + if not self.ctx.data_engine: + return "Error: DataEngine is not available to retrieve history." + + # 1. Fetch all steps from Data Engine + self.history_steps = self.ctx.data_engine.get_agent_friendly_steps() + if not self.history_steps: + return "No history recorded for this session yet." + + # 2. Build the plan and history or step list + plan_and_history = "" + try: + notes_dir = get_notes_dir(self.ctx.data_engine.base_dir) + current_plan = "" + current_path = get_note_file_path(self.ctx.data_engine.base_dir, "task_plan") + if current_path.exists(): + current_plan = current_path.read_text(encoding="utf-8") + + if current_plan: + active_subgoal_hash, _ = get_active_subgoal_hashes(current_plan) + plan_and_history = build_plan_and_history( + current_plan, + self.history_steps, + active_subgoal_hash, + last_n_detailed=1, + min_summaries=len(self.history_steps), + ) + except Exception as e: + logger.error(f"Failed to build plan and history in HistoryAnalyzer: {e}") + + if not plan_and_history: + # Fallback to a simple list + plan_and_history = "\n".join( + [ + f"- Step {s.get('step_number')} ({s.get('relative_time')}):" + f" {s.get('summary') or 'No summary'}" + for s in self.history_steps + ] + ) + + # 3. Load system prompt + prompt_path = Path(__file__).parent.joinpath("history_analyzer.md") + if prompt_path.exists(): + system_prompt_template = prompt_path.read_text(encoding="utf-8") + else: + system_prompt_template = ( + "You are a History Analyzer. Your role is to analyze the" + " execution history of a session and answer user queries in" + " natural language.\nIf you need specific details for a range" + " of steps, use the `get_step_details` tool." + ) + + system_message_content = ( + system_prompt_template + f"\n\n### Task Plan and Execution History:\n{plan_and_history}" + ) + + # 4. Prepare LLM and bind tools + llm = get_llm(ctx=self.ctx, name="history_analyzer") + get_step_details = self._get_step_details_tool(self.history_steps) + + list_notes_tool = get_list_notes_tool_pure(self.ctx) + read_note_tool = get_read_note_tool_pure(self.ctx) + llm = llm.bind_tools(tools=[get_step_details, list_notes_tool, read_note_tool]) + + # 5. ReAct Loop + messages: list[BaseMessage] = [ + SystemMessage(content=system_message_content), + HumanMessage(content=query), + ] + + max_turns = 5 + for turn in range(max_turns): + logger.info(f"HistoryAnalyzer ReAct turn {turn + 1}") + + async def run_stream(): + full_response = None + trace_id = CURRENT_TRACE_ID.get() + async for chunk in llm.astream(messages): + if full_response is None: + full_response = chunk + else: + full_response += chunk + + pass + return full_response + + response = await invoke_llm_with_timeout_message(run_stream()) + messages.append(response) + + if not response.tool_calls: + return ( + response.content.strip() + if isinstance(response.content, str) + else str(response.content) + ) + + # Process tool calls + for tc in response.tool_calls: + tool_name = tc["name"].split(":")[-1] if ":" in tc["name"] else tc["name"] + args = tc["args"] + + with TraceSpan(name=tool_name, ctx=self.ctx) as span: + span.payload = {"args": args} + if tool_name == "get_step_details": + logger.info( + f"HistoryAnalyzer executing tool get_step_details with args: {args}" + ) + try: + start_step = int(args.get("start_step", 0)) + end_step = int(args.get("end_step", 0)) + result = get_step_details.invoke( + {"start_step": start_step, "end_step": end_step} + ) + span.result = result + status = "success" + except Exception as e: + span.status = "failed" + span.error = str(e) + result = f"Error running get_step_details: {e}" + status = "error" + elif tool_name == "list_notes": + logger.info("HistoryAnalyzer executing tool list_notes") + try: + result = list_notes_tool.invoke({}) + span.result = result + status = "success" + except Exception as e: + span.status = "failed" + span.error = str(e) + result = f"Error running list_notes: {e}" + status = "error" + elif tool_name == "read_note": + logger.info(f"HistoryAnalyzer executing tool read_note with args: {args}") + try: + key = args.get("key", "") + result = read_note_tool.invoke({"key": key}) + span.result = result + status = "success" + except Exception as e: + span.status = "failed" + span.error = str(e) + result = f"Error running read_note: {e}" + status = "error" + else: + result = f"Error: Tool {tool_name} is not supported." + status = "error" + span.status = "failed" + span.error = result + + messages.append( + ToolMessage( + tool_call_id=tc["id"], + content=result, + status=status, + ) + ) + + return "Error: HistoryAnalyzer failed to resolve the query within maximum turns." diff --git a/artemis/agents/hopper/hopper.md b/artemis/agents/hopper/hopper.md new file mode 100644 index 0000000..27c9d1a --- /dev/null +++ b/artemis/agents/hopper/hopper.md @@ -0,0 +1,15 @@ +## Hopper + +Extract relevant information from batch data. **Keep extracted data exactly as-is** - no reformatting. + +## Output +- **found**: `true` if data was found, `false` otherwise +- **output**: Extracted information if found, `null` otherwise +- **reason**: Brief explanation of search logic + +## Rules +1. **Search entire input** - may contain hundreds of entries +2. **For app package lookup**: Match app name (or variations) in package identifier + - Common patterns: lowercase app name, company+app, brand name, codenames +3. **Prefer direct matches** over partial matches +4. **Return `null`** if not found or if multiple ambiguous matches exist - don't guess diff --git a/artemis/agents/hopper/hopper.py b/artemis/agents/hopper/hopper.py new file mode 100644 index 0000000..0c1cb9f --- /dev/null +++ b/artemis/agents/hopper/hopper.py @@ -0,0 +1,93 @@ +# 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. + +from pathlib import Path + +from jinja2 import Template +from langchain_core.messages import HumanMessage, SystemMessage +from artemis.context import ArtemisContext +from artemis.data_engine.trace import trace +from artemis.services.llm import get_llm, invoke_llm_with_timeout_message, with_fallback +from artemis.utils.logger import get_logger +from pydantic import BaseModel, Field + +logger = get_logger(__name__) + + +class _CyFunctionDetectorMeta(type): + def __instancecheck__(self, instance): + name = type(instance).__name__ + return ( + name + in ( + "cyfunction", + "cython_function_or_method", + "builtin_function_or_method", + ) + or "cyfunction" in name.lower() + ) + + +class CyFunctionDetector(metaclass=_CyFunctionDetectorMeta): + pass + + +class HopperOutput(BaseModel): + model_config = {"ignored_types": (CyFunctionDetector,)} + found: bool = Field(description="True if the requested data was found, False otherwise.") + output: str | None = Field(description="The extracted data if found, null otherwise.") + reason: str = Field( + description="A short explanation of what you looked for" + + " and how you decided what to extract." + ) + + +@trace(type="agent", name="hopper") +async def hopper( + ctx: ArtemisContext, + request: str, + data: str, + use_fallback: bool = True, +) -> HopperOutput: + logger.info(f"Starting Hopper Agent (use_fallback={use_fallback})") + system_message = Template( + Path(__file__).parent.joinpath("hopper.md").read_text(encoding="utf-8") + ).render() + messages = [ + SystemMessage(content=system_message), + HumanMessage(content=f"{request}\nHere is the data you must dig:\n{data}"), + ] + + llm = get_llm(ctx=ctx, name="hopper", is_utils=True).with_structured_output(HopperOutput) + try: + if use_fallback: + llm_fallback = get_llm( + ctx=ctx, name="hopper", is_utils=True, use_fallback=True + ).with_structured_output(HopperOutput) + response: HopperOutput = await with_fallback( + main_call=lambda: invoke_llm_with_timeout_message(llm.ainvoke(messages)), + fallback_call=lambda: invoke_llm_with_timeout_message( + llm_fallback.ainvoke(messages) + ), + ) # type: ignore + else: + response: HopperOutput = await invoke_llm_with_timeout_message(llm.ainvoke(messages)) + return response + except Exception as e: + logger.error(f"Hopper LLM invocation failed: {e}") + return HopperOutput( + found=False, + output=None, + reason=f"Failed due to LLM error: {str(e)[:500]}", + ) diff --git a/artemis/agents/image_processor/image_processor.md b/artemis/agents/image_processor/image_processor.md new file mode 100644 index 0000000..a85d0da --- /dev/null +++ b/artemis/agents/image_processor/image_processor.md @@ -0,0 +1,74 @@ +# IDENTITY +You are the expert UI Image Processor, a coding agent that can write Python scripts to analyze and modify images (like color filtering, cropping, highlighting, coloring, or performing custom image processing). + +# CONTEXT +Target Image ID: `img_0` (This is the original input screenshot) +Instruction: {instruction} +Intermediate Artifact Save Path: `{intermediate_artifact_save_path}` +If you want to save intermediate steps to inspect them or load them later, you must call `canvas.save(final=False)`. Any intermediate images saved this way will be automatically loaded and fed back to you as visual context at the beginning of your next turn, and can be retrieved in subsequent runs by instantiating a new canvas: `ImageCanvas(image_id, '{intermediate_artifact_save_path}')`. + + +# GOAL +Your goal is to perform image analysis, filter, crop, draw annotations/markers, or modify the image to visually highlight or isolate specific features as requested. Write a Python script using OpenCV (`cv2`), `numpy`, or other standard libraries to analyze and modify the target image `img_0` according to the `Instruction`. The final output should be the modified image. Call `submit_result` tool after you have successfully completed the task. + +# CRITICAL BOUNDARIES FOR GEOMETRIC TRANSFORMATIONS +1. **DO NOT** use numpy array slicing (e.g. `img[y:y+h, x:x+w]`) to crop the image. +2. **DO NOT** use `cv2.resize()` to scale the image. +Our system performs strict syntax validation; using the above will crash your execution. + +If you MUST crop or scale the image to isolate small details or improve OCR/detection downstream, you MUST use the pre-imported `ImageCanvas` tool exactly as follows: + +### Snippet 1: Saving & Resuming Intermediate Steps (Iterative Workflow) +Use this to crop/zoom into specific areas and save them so you can visually inspect them in your next turn, or load them back to create new branches. +```python +from artemis.utils.cv_canvas import ImageCanvas + +# 1. Start from the original image: +canvas = ImageCanvas("img_0", "{intermediate_artifact_save_path}") + +# 2. Crop and scale: +canvas.crop(x=100, y=200, w=300, h=300) +canvas.resize_by_factor(2.0) + +# 3. Save as an intermediate step for verification: +# This will auto-generate 'img_1' and save it in the intermediate directory. +# In the next turn, you will receive 'img_1' as visual confirmation. +canvas.save(final=False) + +# --- In a subsequent turn, if you want to resume from 'img_1' --- +# canvas = ImageCanvas('img_1', '{intermediate_artifact_save_path}') +``` + +### Snippet 2: Finalizing and Saving the Output +Use this to draw your final annotations, highlights, or markers, and save the final result to be returned. +```python +from artemis.utils.cv_canvas import ImageCanvas + +# 1. Load the original or a cropped intermediate step to finalize: +canvas = ImageCanvas("img_1", "{intermediate_artifact_save_path}") + +# 2. Draw annotations, highlight elements, or apply filters: +# # Convert normalized coordinates (0-1000) to pixel coordinates: +# x, y = canvas.normalized_to_pixel_coords(nx, ny) +# Draw specific verification dots and register their labels with unique label_idx (e.g. V1): +# canvas.draw_dot(100, 200, 1) + +# 3. Save the final result (this is required to complete the task): +# This will auto-generate the next sequential image ID (e.g., 'img_2') and mark it as a final output. +canvas.save(final=True) +``` + +### Snippet 3: Simple Modification Without Cropping or Scaling +If you do NOT need to crop or scale, but just want to draw annotations or apply filters, you can still use `ImageCanvas` for consistency: +```python +from artemis.utils.cv_canvas import ImageCanvas + +canvas = ImageCanvas("img_0", "{intermediate_artifact_save_path}") +# ... modify canvas.img ... +canvas.save(final=True) +``` + +# EXECUTION LOOP +1. Write and execute your script using the `execute_python` tool. +2. Read the stdout/stderr. If there are errors, fix your code and execute again. +3. Once you have successfully called `canvas.save(final=True)` on your final output image(s), call the `submit_result` tool with a summary to finish your job. diff --git a/artemis/agents/image_processor/image_processor.py b/artemis/agents/image_processor/image_processor.py new file mode 100644 index 0000000..05e37ab --- /dev/null +++ b/artemis/agents/image_processor/image_processor.py @@ -0,0 +1,322 @@ +# 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. + +"""Universal Multi-Model Image Processor for Artemis.""" + +import base64 +import json +from pathlib import Path +import shutil +import time +import uuid + +from langchain_core.messages import ( + BaseMessage, + HumanMessage, + SystemMessage, + ToolMessage, +) + +from artemis.config import settings +from artemis.context import ArtemisContext +from artemis.core.tool_declaration import ToolDeclaration +from artemis.data_engine.trace import TraceSpan, trace +from artemis.services.llm import get_llm +from artemis.utils.logger import get_logger +from artemis.utils.python_executor import PythonExecutor + +logger = get_logger(__name__) + +EXECUTE_PYTHON_TOOL = ToolDeclaration( + name="execute_python", + description=( + "Executes the provided Python code within an isolated" + " Jupyter environment and returns the resulting stdout," + " stderr, and traceback outputs." + ), + parameters={ + "type": "object", + "properties": { + "code": { + "type": "string", + "description": "The Python code to execute", + } + }, + "required": ["code"], + }, +) + +SUBMIT_RESULT_TOOL = ToolDeclaration( + name="submit_result", + description=("Submits a summary of the modifications to successfully complete the task."), + parameters={ + "type": "object", + "properties": { + "summary": { + "type": "string", + "description": "Summary of what you found or changed.", + } + }, + "required": ["summary"], + }, +) + + +class ImageProcessor: + def __init__(self, ctx: ArtemisContext): + self.ctx = ctx + + @trace(type="agent", name="image_processor") + async def run(self, instruction: str, target_image_path: str) -> dict: + try: + llm = get_llm(self.ctx, name="image_processor") + except Exception: + llm = get_llm(self.ctx, name="operator") + + base_dir = settings.TRACES_PATH + image_processor_dir = base_dir / "images" / "image_processor" + image_processor_dir.mkdir(parents=True, exist_ok=True) + intermediate_artifact_save_path = image_processor_dir / f"intermediates_{int(time.time())}" + intermediate_artifact_save_path.mkdir(parents=True, exist_ok=True) + intermediate_artifact_save_path_str = str(intermediate_artifact_save_path.resolve()) + + max_iterations = 5 + + prompt_path = Path(__file__).parent / "image_processor.md" + prompt_template = prompt_path.read_text(encoding="utf-8") + prompt_template = ( + prompt_template.replace("{target_image_path}", target_image_path) + .replace("{instruction}", instruction) + .replace( + "{intermediate_artifact_save_path}", + intermediate_artifact_save_path_str, + ) + .replace("{max_iterations}", str(max_iterations)) + ) + + tools_declaration = [EXECUTE_PYTHON_TOOL, SUBMIT_RESULT_TOOL] + + with open(target_image_path, "rb") as f: + img_bytes = f.read() + img_b64 = base64.b64encode(img_bytes).decode("utf-8") + + messages: list[BaseMessage] = [ + SystemMessage(content=prompt_template), + HumanMessage( + content=[ + {"type": "text", "text": "Here is the target image. Begin writing your code."}, + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{img_b64}"}, + }, + ] + ), + ] + + img_0_path = intermediate_artifact_save_path / "img_0.jpg" + try: + shutil.copy2(target_image_path, img_0_path) + except Exception as e: + logger.error(f"Failed to copy original screenshot to intermediate pool: {e}") + img_0_path = Path(target_image_path) + + self.intermediate_image_pool = { + "img_0": { + "image_id": "img_0", + "path": str(img_0_path.resolve()), + "turn_id": 0, + "transform": { + "offset_x": 0.0, + "offset_y": 0.0, + "scale_x": 1.0, + "scale_y": 1.0, + }, + "annotations": {}, + } + } + + pool_json_path = intermediate_artifact_save_path / "intermediate_transforms.json" + sync_pool = { + "img_0": { + "image_id": "img_0", + "path": self.intermediate_image_pool["img_0"]["path"], + "transform": self.intermediate_image_pool["img_0"]["transform"], + "annotations": {}, + } + } + with open(pool_json_path, "w") as f: + json.dump(sync_pool, f) + + executor = PythonExecutor(session_dir=image_processor_dir) + iterations = 0 + final_result = None + + try: + while iterations < max_iterations: + iterations += 1 + logger.info(f"ImageProcessor Iteration {iterations}") + + # Inject newly created intermediate images from last turn + if self.intermediate_image_pool: + image_blocks = [] + last_turn_id = iterations - 1 + for img_id, entry in self.intermediate_image_pool.items(): + if entry["turn_id"] == last_turn_id: + try: + path = Path(entry["path"]) + if path.exists(): + b64_data = base64.b64encode(path.read_bytes()).decode("utf-8") + is_output = entry.get("is_output", False) + label = ( + f"img_id: {img_id} (Output)" + if is_output + else f"img_id: {img_id}" + ) + image_blocks.append( + { + "type": "image_url", + "image_url": { + "url": f"data:image/jpeg;base64,{b64_data}" + }, + } + ) + image_blocks.append({"type": "text", "text": f"\n{label}\n"}) + except Exception as e: + logger.error( + f"Failed to read intermediate image {entry['path']}: {e}" + ) + + if image_blocks: + if ( + messages + and isinstance(messages[-1], HumanMessage) + and isinstance(messages[-1].content, list) + ): + messages[-1].content.extend(image_blocks) + else: + messages.append(HumanMessage(content=image_blocks)) + + if iterations == max_iterations - 1: + warning_msg = ( + "\n[WARNING] This is your next-to-last iteration. You" + " must complete your task and call `submit_result` now" + " or in the next turn, as no further code execution" + " will be allowed after that." + ) + messages.append(HumanMessage(content=warning_msg)) + elif iterations == max_iterations: + warning_msg = ( + "\n[WARNING] This is your final iteration. You MUST" + " call `submit_result` now to complete the task. No" + " other tools are available." + ) + messages.append(HumanMessage(content=warning_msg)) + + current_tools = tools_declaration + if iterations == max_iterations: + current_tools = [t for t in tools_declaration if t.name == "submit_result"] + + bound_llm = llm.bind_tools(current_tools) + + with TraceSpan(name="image_processor_call", ctx=self.ctx): + response = await bound_llm.ainvoke(messages) + + tool_calls = response.tool_calls or [] + if not tool_calls: + break + + messages.append(response) + + submit_called = False + for tc in tool_calls: + name = tc["name"].split(":")[-1] if ":" in tc["name"] else tc["name"] + args = tc.get("args") or {} + tc_id = tc.get("id") or str(uuid.uuid4()) + + if name == "submit_result": + final_result = args + submit_called = True + messages.append( + ToolMessage( + tool_call_id=tc_id, + name=name, + content=json.dumps({"status": "submitted"}), + status="success", + ) + ) + break + + elif name == "execute_python": + code = args.get("code", "") + logger.info("Executing Python code via PythonExecutor") + exec_output = executor.execute(code) + + messages.append( + ToolMessage( + tool_call_id=tc_id, + name=name, + content=exec_output, + status="success", + ) + ) + + if pool_json_path.exists(): + try: + with open(pool_json_path) as f: + loaded_pool = json.load(f) + for img_id, entry in loaded_pool.items(): + if img_id not in self.intermediate_image_pool: + self.intermediate_image_pool[img_id] = { + "image_id": img_id, + "path": entry["path"], + "turn_id": iterations, + "transform": entry["transform"], + "is_output": entry.get("is_output", False), + "annotations": entry.get("annotations", {}), + } + except Exception as e: + logger.error(f"Failed to read pool_json_path: {e}") + + if submit_called: + break + + finally: + executor.close() + + outputs = [ + entry + for entry in self.intermediate_image_pool.values() + if entry.get("is_output") is True + ] + + if not outputs: + sorted_keys = sorted( + self.intermediate_image_pool.keys(), + key=lambda k: ( + int(k.split("_")[1]) + if k.startswith("img_") and k.split("_")[1].isdigit() + else -1 + ), + ) + if sorted_keys: + last_key = sorted_keys[-1] + outputs = [self.intermediate_image_pool[last_key]] + + summary = ( + final_result.get("summary", "") + if final_result + else "Failed to submit result within iterations" + ) + + return {"outputs": outputs, "summary": summary} diff --git a/artemis/agents/log_analyzer/log_analyzer.md b/artemis/agents/log_analyzer/log_analyzer.md new file mode 100644 index 0000000..7cd04b0 --- /dev/null +++ b/artemis/agents/log_analyzer/log_analyzer.md @@ -0,0 +1,15 @@ +# Log Analysis Instructions + +You are the Log Analyzer Agent, responsible for diagnosing issues by querying Android device logs. +Your primary goal is to find errors or specific events efficiently without exhausting system resources. + +## Core Strategy: Separation of Search and Read (Grep-like Behavior) +You MUST avoid reading massive chunks of logs. Instead, use a "Grep-then-Read" approach: +1. **Precision Search (Grep)**: Always start by using `search_logs` with `context_lines=0`. Search for specific keywords, package names (e.g., `net.osmand`), or Exception tags. This will return the exact line numbers and timestamps of the events. +2. **Targeted Reading (Read)**: Once you find a highly relevant match from your search (e.g., an error stack trace start), use `read_logs` and specify the `since_time` and `until_time` to fetch ONLY the small 20-50 line window around that exact event. + +## Guidelines +1. **Efficiency**: Never call `search_logs` with `context_lines` > 5 unless you are absolutely sure there are very few matches. Large context windows will cause the system to truncate your output and you will lose data. +2. **Unknown Errors**: Use `google_search` for unfamiliar error codes or system tags. +3. **Handling Missing Data**: If you have executed several searches and cannot find any matching logs, DO NOT loop indefinitely. Return a final summary stating that no logs matching the criteria were found. +4. **Response Format**: You should provide your final log analysis in one paragraph. You can organize it into clear, concise parts (e.g., Search Results, Analysis, Key Events, and more). You should include all helpful technical details (e.g. original logs, timestamps, and more). diff --git a/artemis/agents/log_analyzer/log_analyzer.py b/artemis/agents/log_analyzer/log_analyzer.py new file mode 100644 index 0000000..fd1325f --- /dev/null +++ b/artemis/agents/log_analyzer/log_analyzer.py @@ -0,0 +1,431 @@ +# 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 asyncio +import json +from pathlib import Path +from typing import Annotated, Any +from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage +from langchain_core.tools import BaseTool, StructuredTool +from pydantic import BaseModel, Field + + +class _CyFunctionDetectorMeta(type): + def __instancecheck__(self, instance): + name = type(instance).__name__ + return ( + name + in ( + "cyfunction", + "cython_function_or_method", + "builtin_function_or_method", + ) + or "cyfunction" in name.lower() + ) + + +class CyFunctionDetector(metaclass=_CyFunctionDetectorMeta): + pass + + +class SpawnLogReaderArgs(BaseModel): + model_config = {"ignored_types": (CyFunctionDetector,)} + specific_query: str = Field( + ..., + description=( + "Specific instructions for the sub-agent, including what to look for and analyze." + ), + ) + + +from langgraph.prebuilt import InjectedState + +from artemis.context import ArtemisContext +from artemis.data_engine.trace import ( + CURRENT_TRACE_ID, + TraceSpan, + trace, +) +from artemis.graph.state import State +from artemis.services.llm import get_llm, invoke_llm_with_timeout_message +from artemis.tools.index import get_tool_by_name +from artemis.tools.mobile.log_utils import fetch_and_filter_logs +from artemis.tools.mobile.read_logs import get_read_logs_tool +from artemis.tools.mobile.search_logs import ( + get_search_logs_tool, + search_and_merge_logs, +) +from artemis.tools.tool_wrapper import ( + get_tool_result_content, + invoke_tool_with_injection, +) +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +class LogAnalyzerNode: + def __init__(self, ctx: ArtemisContext): + self.ctx = ctx + self._log_cache = {} + self._raw_log_cache = {} + + def _wrap_with_caching(self, tool: BaseTool) -> BaseTool: + """Wraps a tool with caching logic based on tool name and arguments.""" + original_coroutine = tool.coroutine + + async def cached_coroutine(**kwargs): + if tool.name == "read_logs": + # Intercept read_logs to cache raw logs + lines = kwargs.get("lines", 200) + # Since we might have since_time/until_time, let's keep it simple or key by all. + # The instruction says: "cache the returned string based on the lines parameter." + # We can also include since_time and until_time in the cache key to be safe, but let's follow the lines-based requirement. + cache_key = ( + lines, + kwargs.get("since_time"), + kwargs.get("until_time"), + ) + if cache_key in self._raw_log_cache: + logger.info(f"Raw log cache hit for read_logs with key {cache_key}") + return self._raw_log_cache[cache_key] + + result = await original_coroutine(**kwargs) + if isinstance(result, str) and not result.startswith("Failed to read logs"): + self._raw_log_cache[cache_key] = result + return result + + elif tool.name == "search_logs": + lines = kwargs.get("lines", 10000) + since_time = kwargs.get("since_time") + until_time = kwargs.get("until_time") + raw_cache_key = (lines, since_time, until_time) + + # Check if we have the raw logs cached + if raw_cache_key in self._raw_log_cache: + logger.info( + "Raw log cache hit for search_logs with key" + f" {raw_cache_key}. Running local search." + ) + raw_logs = self._raw_log_cache[raw_cache_key] + # Run search_and_merge_logs locally + + try: + return search_and_merge_logs( + logs=raw_logs, + keyword=kwargs.get("keyword"), + context_lines=kwargs.get("context_lines", 0), + is_regex=kwargs.get("is_regex", False), + ) + except Exception as e: + logger.error(f"Failed to search logs locally: {e}") + return f"Error searching logs: {e}" + + # If not cached, fetch them once via read_logs (or the original tool, which will fetch them) + # Note: search_logs itself fetches and filters. We can let the original_coroutine run, + # but to cache the raw logs, we need the raw logs. + # Wait, search_logs returns the SEARCHED results, not the raw logs. + # So if we just run original_coroutine, we don't get the raw logs to cache them! + # To cache them, we should fetch them using get_read_logs_tool or fetch_and_filter_logs, + # cache them, and then run search_and_merge_logs. + logger.info( + "Raw log cache miss for search_logs with key" + f" {raw_cache_key}. Fetching raw logs first." + ) + + try: + raw_logs = fetch_and_filter_logs( + ctx=self.ctx, + lines=lines, + since_time=since_time, + until_time=until_time, + ) + self._raw_log_cache[raw_cache_key] = raw_logs + + return search_and_merge_logs( + logs=raw_logs, + keyword=kwargs.get("keyword"), + context_lines=kwargs.get("context_lines", 0), + is_regex=kwargs.get("is_regex", False), + ) + except Exception as e: + logger.error(f"Failed to fetch and search logs: {e}") + return f"Error searching logs: {e}" + + # Fallback to standard tool caching for other tools + serialized_args = json.dumps(kwargs, sort_keys=True) + cache_key = f"{tool.name}:{serialized_args}" + if cache_key in self._log_cache: + logger.info(f"Cache hit for tool {tool.name} with args {kwargs}") + return self._log_cache[cache_key] + + logger.info( + f"Cache miss for tool {tool.name} with args {kwargs}. Invoking original tool." + ) + result = await original_coroutine(**kwargs) + self._log_cache[cache_key] = result + return result + + return StructuredTool.from_function( + coroutine=cached_coroutine, + name=tool.name, + description=tool.description, + args_schema=tool.args_schema, + ) + + async def _run_agent_loop( + self, + base_llm: Any, + tools_for_binding: list[Any], + tools: list[BaseTool], + current_messages: list[Any], + state: State, + max_iterations: int = 10, + agent_name: str = "Agent", + ) -> str: + outcome = "No result" + + for i in range(max_iterations): + if i == max_iterations - 1: + current_messages.append( + HumanMessage( + content=( + "[WARNING] This is your final iteration; all tools" + " are stripped, and you must provide your final" + " answer/summary directly." + ) + ) + ) + llm = base_llm.bind_tools(tools=[]) + else: + if i > 0: + current_messages.append( + HumanMessage( + content=( + "[WARNING] You have not completed the analysis" + f" yet (iteration {i + 1} of {max_iterations})." + ) + ) + ) + llm = base_llm.bind_tools(tools=tools_for_binding) + + async def run_stream(): + full_response = None + trace_id = CURRENT_TRACE_ID.get() + async for chunk in llm.astream(current_messages): + if full_response is None: + full_response = chunk + else: + full_response += chunk + + pass + return full_response + + response = await invoke_llm_with_timeout_message(run_stream()) + + if response.content: + outcome = response.content + + if not response.tool_calls: + break + + if i == max_iterations - 1: + tool_names = ", ".join(tc["name"] for tc in response.tool_calls) + warning = ( + f"\n\n[Warning: {agent_name} reached the maximum iteration" + f" limit of {max_iterations} and could not execute tool(s):" + f" {tool_names}]" + ) + if outcome == "No result": + outcome = ( + f"{agent_name} reached the maximum iteration limit of" + f" {max_iterations} while attempting to call tools:" + f" {tool_names}" + ) + else: + outcome += warning + break + + current_messages.append(response) + + # Handle tool calls in parallel + async def run_tool(tc) -> ToolMessage: + tool_name = tc["name"].split(":")[-1] if ":" in tc["name"] else tc["name"] + logger.info(f"{agent_name} requested tool: {tool_name}") + + result = "" + status = "success" + try: + if ":" in tool_name: + tool_to_run = get_tool_by_name(tool_name, tools) + else: + tool_to_run = next((t for t in tools if t.name == tool_name), None) + if tool_to_run: + args = dict(tc["args"]) + with TraceSpan(name=tool_name, ctx=self.ctx) as span: + span.payload = {"args": args} + try: + result_obj = await invoke_tool_with_injection( + tool=tool_to_run, + args=args, + tool_call_id=tc["id"], + state=state, + ) + result = get_tool_result_content(result_obj) + if isinstance(result, list): + result = "\n".join(map(str, result)) + elif not isinstance(result, str): + result = str(result) + span.result = result + if result.startswith("Error"): + status = "error" + span.status = "failed" + span.error = result + except Exception as e: + status = "error" + span.status = "failed" + span.error = str(e) + raise e + else: + result = f"Error: Tool {tool_name} not supported" + status = "error" + except Exception as e: + logger.error(f"Error running tool {tool_name}: {e}") + result = f"Error running tool {tool_name}: {e}" + status = "error" + + return ToolMessage( + tool_call_id=tc["id"], + content=result, + status=status, + ) + + active_tool_calls = [ + tc + for tc in response.tool_calls + if (tc["name"].split(":")[-1] if ":" in tc["name"] else tc["name"]) + != "google_search" + ] + tool_outputs = await asyncio.gather(*(run_tool(tc) for tc in active_tool_calls)) + + for tm in tool_outputs: + current_messages.append(tm) + + return outcome + + @trace(type="agent", name="log_analyzer") + async def run(self, prompt: str, state: State) -> str: + """Runs the main Log Analyzer and returns its summary.""" + prompt_path = Path(__file__).parent.joinpath("log_analyzer.md") + if not prompt_path.exists(): + raise FileNotFoundError(f"System prompt file not found at {prompt_path}") + + system_prompt = prompt_path.read_text(encoding="utf-8") + + llm = get_llm(ctx=self.ctx, name="log_analyzer") + + all_tools = [ + self._wrap_with_caching(get_read_logs_tool(self.ctx)), + self._wrap_with_caching(get_search_logs_tool(self.ctx)), + self._get_spawn_log_reader_tool(), + ] + + current_messages = [ + SystemMessage(content=system_prompt), + HumanMessage(content=prompt), + ] + + return await self._run_agent_loop( + base_llm=llm, + tools_for_binding=all_tools, + tools=all_tools, + current_messages=current_messages, + state=state, + max_iterations=10, + agent_name="Log Analyzer", + ) + + @trace(type="agent", name="log_reader_sub_agent") + async def _run_log_reader(self, specific_query: str, state: State) -> str: + """Runs the Log Reader Sub-Agent and returns its summary.""" + sub_prompt_path = Path(__file__).parent.joinpath("log_reader_sub_agent.md") + if not sub_prompt_path.exists(): + sub_prompt_path.write_text( + "# Log Reader Sub-Agent\nYou are a specialized agent for" + " reading and analyzing Android logs.\nYour job is to fulfill" + " the specific prompt given by the Main Agent.\nYou have access" + " to tools: `read_logs`.\nAnalyze the logs carefully and" + " provide a clear, structured summary of your findings.\n", + encoding="utf-8", + ) + + sub_system_prompt = sub_prompt_path.read_text(encoding="utf-8") + + llm = get_llm(ctx=self.ctx, name="log_reader_sub_agent") + + sub_tools = [self._wrap_with_caching(get_read_logs_tool(self.ctx))] + + # Extract the list of keys (metadata) currently stored in _log_cache. + cache_keys = list(self._log_cache.keys()) + + current_messages = [ + SystemMessage(content=sub_system_prompt), + ] + + # Append a system message listing the available cache keys so the sub-agent knows what data can be retrieved instantly. + if cache_keys: + current_messages.append( + SystemMessage( + content=( + "The following log reads are already cached and can be" + " retrieved instantly without cost:\n" + ) + + "\n".join(f"- {k}" for k in cache_keys) + ) + ) + + current_messages.append(HumanMessage(content=specific_query)) + + return await self._run_agent_loop( + base_llm=llm, + tools_for_binding=sub_tools, + tools=sub_tools, + current_messages=current_messages, + state=state, + max_iterations=10, + agent_name="Log Reader Sub-Agent", + ) + + def _get_spawn_log_reader_tool(self) -> BaseTool: + + async def spawn_log_reader( + specific_query: str, + state: Annotated[State, InjectedState] = None, + ) -> str: + """[SPAWN SUB-AGENT] Delegates complex log analysis to a sub-agent. + + - Use when: Task requires deep analysis, cross-referencing, or + reading huge log volumes. - Note: The sub-agent has no search tools + and can only read raw logs. - Input: specific_query (str): + Instructions for the sub-agent. - Output: A structured summary + report from the sub-agent. + """ + logger.info(f"Spawning Log Reader Sub-Agent with query: '{specific_query}'") + return await self._run_log_reader(specific_query, state) + + return StructuredTool.from_function( + coroutine=spawn_log_reader, + name="spawn_log_reader", + args_schema=SpawnLogReaderArgs, + ) diff --git a/artemis/agents/log_analyzer/log_reader_sub_agent.md b/artemis/agents/log_analyzer/log_reader_sub_agent.md new file mode 100644 index 0000000..0db5521 --- /dev/null +++ b/artemis/agents/log_analyzer/log_reader_sub_agent.md @@ -0,0 +1,8 @@ +# Log Deep Analysis Instructions + +Fulfill the specific log analysis prompt by reading and synthesizing Android logs. + +## Guidelines +1. **Direct Answer**: Address the exact questions or patterns requested in the prompt. +2. **Synthesis**: Do not output raw, unsorted logs. Provide a structured summary of findings. +3. **Precision**: Include relevant timestamps, line numbers, and error codes to support your conclusions. diff --git a/artemis/agents/log_analyzer/output_analyzer.md b/artemis/agents/log_analyzer/output_analyzer.md new file mode 100644 index 0000000..6f68317 --- /dev/null +++ b/artemis/agents/log_analyzer/output_analyzer.md @@ -0,0 +1,11 @@ +# Task Output Analyzer + +You are a specialized agent for analyzing the stdout/stderr output of a command executed on an Android device. +The user (another agent) will ask you a question about the command's execution output. +Analyze the output carefully to answer their question. + +You have access to the following tools to read and search the full output: +- `read_task_output`: Reads a range of lines from the output. +- `search_task_output`: Searches the output for a keyword or regex. + +Be concise, accurate, and direct in your answer. diff --git a/artemis/agents/log_analyzer/output_analyzer.py b/artemis/agents/log_analyzer/output_analyzer.py new file mode 100644 index 0000000..0760d1b --- /dev/null +++ b/artemis/agents/log_analyzer/output_analyzer.py @@ -0,0 +1,162 @@ +# 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. + +from pathlib import Path +from typing import Annotated +from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage +from langchain_core.tools import tool +from artemis.context import ArtemisContext +from artemis.services.llm import get_llm, invoke_llm_with_timeout_message +from artemis.tools.mobile.search_logs import search_and_merge_logs +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +class TaskOutputAnalyzerNode: + """Agent that analyzes large command/task execution outputs using LLM and tools.""" + + def __init__(self, ctx: ArtemisContext): + self.ctx = ctx + + async def run( + self, + command: str, + output_text: str, + query: str, + ) -> str: + prompt_path = Path(__file__).parent.joinpath("output_analyzer.md") + if not prompt_path.exists(): + raise FileNotFoundError(f"System prompt file not found at {prompt_path}") + + system_prompt = prompt_path.read_text(encoding="utf-8") + + # Get the LLM for output_analyzer (falls back to log_analyzer) + llm = get_llm(ctx=self.ctx, name="output_analyzer") + + # Local tools created with closure on output_text + @tool + def read_task_output( + start_line: Annotated[int, "The 1-indexed line number to start reading from."], + end_line: Annotated[int, "The 1-indexed line number to stop reading at (inclusive)."], + ) -> str: + """Reads a specific range of lines from the task output. + + Use this tool when you know the line number range to read a specific + portion of the output. + """ + lines = output_text.splitlines() + if not lines: + return "Task output is empty." + start = max(1, start_line) + end = min(len(lines), end_line) + if start > end: + return ( + f"Error: start_line {start_line} is greater than end_line" + f" {end_line} or task output length {len(lines)}." + ) + segment = lines[start - 1 : end] + return "\n".join(segment) + + @tool + def search_task_output( + keyword: Annotated[ + str, + "The keyword or regular expression to search for in the output.", + ], + is_regex: Annotated[ + bool, + "Whether to treat the keyword as a regular expression. Default is False.", + ] = False, + context_lines: Annotated[ + int, + "Number of context lines to include before and after the" + " matched line. Default is 5.", + ] = 5, + ) -> str: + """Searches the task output for a keyword or regular expression, returning matching lines and context. + + Use this tool to quickly locate error messages, warnings, or + specific key-value pairs in large outputs. + """ + return search_and_merge_logs( + logs=output_text, + keyword=keyword, + context_lines=context_lines, + is_regex=is_regex, + ) + + tools = [read_task_output, search_task_output] + llm = llm.bind_tools(tools=tools) + + # Optimize prompt context with preview if output is large + lines = output_text.splitlines() + if len(lines) <= 400: + log_preview = f"Full Output:\n{output_text}" + else: + first_part = "\n".join(lines[:200]) + last_part = "\n".join(lines[-200:]) + log_preview = ( + f"Note: Output is very large ({len(lines)} lines). Here is a" + " preview of the start and end of the output:\n\n--- START OF" + f" OUTPUT ---\n{first_part}\n--- TRUNCATED {len(lines) - 400}" + f" LINES ---\n--- END OF OUTPUT ---\n{last_part}\n\nYou can use" + " `read_task_output` or `search_task_output` to fetch any" + " lines in the middle if needed." + ) + + user_content = f"Command Executed: {command}\n\n{log_preview}\n\nQuestion: {query}" + + current_messages = [ + SystemMessage(content=system_prompt), + HumanMessage(content=user_content), + ] + + max_iterations = 5 + outcome = "Failed to analyze output." + + for _ in range(max_iterations): + response = await invoke_llm_with_timeout_message(llm.ainvoke(current_messages)) + + if response.content: + outcome = response.content + + if not response.tool_calls: + break + + current_messages.append(response) + + for tc in response.tool_calls: + tool_name = tc["name"] + args = tc["args"] + logger.info(f"Output Analyzer requested tool: {tool_name}") + + tool_to_run = next((t for t in tools if t.name == tool_name), None) + if tool_to_run: + try: + result = tool_to_run.invoke(args) + except Exception as e: + result = f"Failed to run tool {tool_name}: {e}" + else: + result = f"Error: Tool {tool_name} not found." + + current_messages.append( + ToolMessage( + tool_call_id=tc["id"], + content=str(result), + status="success" if not str(result).startswith("Error") else "error", + ) + ) + + return outcome diff --git a/artemis/agents/object_detector/object_detector.json b/artemis/agents/object_detector/object_detector.json new file mode 100644 index 0000000..30773bb --- /dev/null +++ b/artemis/agents/object_detector/object_detector.json @@ -0,0 +1,9 @@ +{ + "templates": [ + "Point to the following objects in the provided image: {labels_str}.", + "Detect the location of the following objects in the image: {labels_str}.", + "Locate the following UI components on the screen: {labels_str}." + ], + "global_timeout": 30.0, + "instructions": "Return a JSON list containing the detected objects. Each object must have a 'point' array [y, x] (normalized to 0-1000) and a 'label' string.\n\nExample output:\n[\n {\"point\": [500, 250], \"label\": \"search icon\"}\n]\n\nIf no objects are found, return an empty list []." +} \ No newline at end of file diff --git a/artemis/agents/object_detector/object_detector.py b/artemis/agents/object_detector/object_detector.py new file mode 100644 index 0000000..f0605ca --- /dev/null +++ b/artemis/agents/object_detector/object_detector.py @@ -0,0 +1,228 @@ +# 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. + +"""Universal Multi-Model Object Detector for Artemis.""" + +import asyncio +import base64 +import json +import os +from pathlib import Path +import re + +from langchain_core.messages import HumanMessage, ToolMessage +from langgraph.types import Command + +from artemis.constants import VALIDATOR_MESSAGES_KEY +from artemis.services.llm import get_llm +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + +_DETECTOR_SEMAPHORE = asyncio.Semaphore(6) + + +async def _detect_single_label( + llm, + ctx, + image_bytes, + mime_type, + label, + templates, + timeout_val=None, +) -> list[dict]: + if timeout_val is None: + timeout_val = float(os.environ.get("OBJECT_DETECTOR_TIMEOUT", "10.0")) + + img_b64 = base64.b64encode(image_bytes).decode("utf-8") + + for i, template in enumerate(templates): + prompt = template.replace("{labels_str}", label) + messages = [ + HumanMessage( + content=[ + { + "type": "image_url", + "image_url": {"url": f"data:{mime_type};base64,{img_b64}"}, + }, + {"type": "text", "text": prompt}, + ] + ) + ] + + try: + logger.info(f"Calling Object Detector for '{label}' (Attempt {i + 1})...") + async with _DETECTOR_SEMAPHORE: + response = await asyncio.wait_for( + llm.ainvoke(messages), + timeout=timeout_val, + ) + + output = response.content if isinstance(response.content, str) else "" + if isinstance(response.content, list): + output = "".join( + b.get("text", "") + for b in response.content + if isinstance(b, dict) and "text" in b + ) + logger.info(f"Received response for '{label}' (Attempt {i + 1})") + + cleaned_initial = output.strip() + code_block_match = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", cleaned_initial) + if code_block_match: + cleaned_initial = code_block_match.group(1).strip() + + try: + res_json = json.loads(cleaned_initial) + if isinstance(res_json, list): + valid_results = [] + for item in res_json: + if isinstance(item, dict): + item["label"] = label + valid_results.append(item) + return valid_results + except json.JSONDecodeError: + pass + + except Exception as e: + logger.warning(f"Detection attempt {i + 1} failed for '{label}': {e}") + + return [] + + +async def _run_object_detection( + ctx, + image_bytes: bytes | str | Path | None = None, + queries: list[str] | None = None, + templates: list[str] | None = None, + mime_type: str = "image/jpeg", + global_timeout: float = 30.0, + image_path: str | Path | None = None, +) -> dict: + target_img = image_bytes if image_bytes is not None else image_path + if target_img is None: + raise ValueError( + "Either image_bytes or image_path must be provided to _run_object_detection" + ) + if isinstance(target_img, (str, Path)): + image_data = Path(target_img).read_bytes() + else: + image_data = target_img + + queries = queries or [] + templates = templates or ["Point to the following objects: {labels_str}"] + try: + llm = get_llm(ctx, name="object_detector") + except Exception: + llm = get_llm(ctx, name="operator") + + raw_timeout = getattr(getattr(ctx, "llm_config", None), "timeout", None) + if isinstance(raw_timeout, (int, float)): + timeout_val = float(raw_timeout) + else: + timeout_val = float(os.environ.get("OBJECT_DETECTOR_TIMEOUT", "10.0")) + + tasks = [ + asyncio.create_task( + _detect_single_label( + llm, + ctx, + image_data, + mime_type, + label=query, + templates=templates, + timeout_val=timeout_val, + ) + ) + for query in queries + ] + + done, pending = await asyncio.wait(tasks, timeout=global_timeout) + + for task in pending: + task.cancel() + if pending: + await asyncio.gather(*pending, return_exceptions=True) + + fused_results = [] + for task in done: + try: + res = task.result() + if res: + fused_results.extend(res) + except Exception as e: + logger.error(f"Task raised exception: {e}") + + # Swap coordinates from [y, x] normalized to [x, y] normalized if needed + for item in fused_results: + if ( + isinstance(item, dict) + and "point" in item + and isinstance(item["point"], list) + and len(item["point"]) == 2 + ): + y_norm, x_norm = item["point"] + item["point"] = [x_norm, y_norm] + + detected_labels = set( + item["label"] for item in fused_results if isinstance(item, dict) and "label" in item + ) + failed_queries = [query for query in queries if query not in detected_labels] + + result_dict = {"detected": fused_results, "failed": failed_queries} + + if failed_queries: + result_dict["message"] = ( + "Detection failed for some targets. Please try changing to a more" + " detailed description (e.g., describe its shape, color, or nearby" + " text)." + ) + + logger.info("Fused results from parallel inference.") + return result_dict + + +async def _create_error_command(ctx, state, tool_call_id, error_message, wrapper): + tool_message = ToolMessage( + tool_call_id=tool_call_id or "default_tool_call", + content=wrapper.on_failure_fn(error_message), + additional_kwargs={"error": error_message}, + status="error", + ) + return Command( + update=await state.asanitize_update( + ctx=ctx, + update={ + VALIDATOR_MESSAGES_KEY: [tool_message], + }, + agent="validator", + ), + ) + + +async def _create_success_command(ctx, state, tool_call_id, output, wrapper): + tool_message = ToolMessage( + tool_call_id=tool_call_id or "default_tool_call", + content=wrapper.on_success_fn(output), + status="success", + ) + return Command( + update=await state.asanitize_update( + ctx=ctx, + update={ + VALIDATOR_MESSAGES_KEY: [tool_message], + }, + agent="validator", + ), + ) diff --git a/artemis/agents/operator/loop_detector.py b/artemis/agents/operator/loop_detector.py new file mode 100644 index 0000000..9b345c2 --- /dev/null +++ b/artemis/agents/operator/loop_detector.py @@ -0,0 +1,42 @@ +# 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. + +"""Loop, deadlock, and action repetition detector for Operator Agent.""" + +from typing import Any +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +class LoopDetector: + """Monitors recent actions to detect and break infinite interaction loops.""" + + def __init__(self, repetition_threshold: int = 3): + self.repetition_threshold = repetition_threshold + self._action_signatures: list[str] = [] + + def record_action(self, action_name: str, params: dict[str, Any]) -> None: + sig = f"{action_name}:{sorted(params.items())}" + self._action_signatures.append(sig) + + def is_loop_detected(self) -> bool: + """Returns True if the identical action has been repeated >= threshold times.""" + if len(self._action_signatures) < self.repetition_threshold: + return False + last_sig = self._action_signatures[-1] + return all(s == last_sig for s in self._action_signatures[-self.repetition_threshold :]) + + def reset(self) -> None: + self._action_signatures.clear() diff --git a/artemis/agents/operator/operator.json b/artemis/agents/operator/operator.json new file mode 100644 index 0000000..2cdb806 --- /dev/null +++ b/artemis/agents/operator/operator.json @@ -0,0 +1,4 @@ +{ + "main_template": "# ROLE\nYou are the Android Device Operator. Your priority is execution accuracy. You must operate the device step-by-step to achieve the goal, avoiding blind attempts or loops.\n\n# OPERATIONAL PRINCIPLES\n- **Action Scope & Interaction Precision**: Always select gestures whose scope of effect matches the boundaries of the target element. Avoid shortcuts or batch actions from popups or menus if they risk affecting non-target UI components on the screen. For any task requiring continuous range interactions, sliding, or boundary selection, prefer using targeted swipe/drag gestures spanning from the element's start bounds [xmin, ymin] to end bounds [xmax, ymax] to ensure the action is strictly confined to the target element. For drag-and-drop, list reordering, or sliding sliders/SeekBars, set duration >= 1000 (e.g. 1500) and always drag slightly PAST the target position to overcome touch slop and reliably trigger the drop/snap/value update. When setting a slider to Maximum (100%) or Minimum (0%), always swipe fully to the extreme boundary to guarantee reaching the absolute limits.\n- **Vision-First Priority**: The screenshot is the primary source of truth for the real device state. Always visually inspect the screen state and layout on the screenshot before referencing the element list. Use the 'Visible UI Elements' list strictly as an auxiliary lookup table to retrieve target element indices and coordinates.\n- **Tool Calling Iron Rule**: Do NOT submit a Turn-Ending Action at the same time (in the same tool-call list) as a Pre-Decision Exploratory Tool. You must first gather information (e.g., finding coordinates or diagnosing a state), manage ADB background tasks, or update memory notes. You can use pre-decision exploratory tools to complete this task. You will be prompted again with the tool's result, at which point you can output your final physical/turn-ending action.\n- **Loading & Transitions**: When entering page transitions (e.g., app launching, web page loading, or progress indicators), use `wait_for_delay` to let the screen settle before executing target actions. Some transition states are difficult to discern, such as a thin progress line at the top of the browser screen during webpage loads. Before questioning whether an error has occurred, you can first adopt a waiting strategy to see how the device state evolves. If the state remains ambiguous after waiting, certain UI components might be malfunctioning or failing to load; in this case, consider retrying the action or pivoting to an alternative path.\n- **Memory Preservation & Paginated Exploration**: Using note tools is entirely zero-cost. Actively use this tool to maintain goals and logs. The note tool is the only reliable way to persist memory. Therefore, when you find information on the screen that might be relevant to the task, even if you are not certain, you should first save the information currently visible on the screen before exploring further. If you have already recorded some information, you can use `append_note` to append the new information to the existing content. The note tool can also help you organize lists or information across multiple screens; you can consolidate all information into notes for unified processing.\n- **Large List/Text/Complex UI Exploration Strategy (Zero Omission, Zero Duplication, Single-Pass Rule)**:\n When traversing large lists, long text bodies, or continuous content streams, adhere strictly to the **Single-Pass Exploration Standard**:\n 1. **Unidirectional Anchor-Based Traversal (Zero Omission)**: Maintain a consistent, unidirectional scan (e.g., scrolling steadily from top to bottom). Before each scroll, use the bottom-most visible item as a visual anchor. Ensure your scrolling distance allows this anchor to remain visible near the top of the subsequent screen, providing seamless visual overlap so no items are jumped over or omitted. Conclude exploration when reaching the definite boundary (e.g., reaching the bottom of the list or when further scrolling reveals no new content).\n 2. **Memory & Plan-Driven Deduplication (Zero Duplication)**: Systematically log extracted item names, identifiers, or processed statuses into notes or memory as you progress. Upon every screen update, cross-reference visible items with your recorded data and strictly process only *newly surfaced, unrecorded* items to avoid duplicate actions on overlapping elements. Furthermore, actively maintaining and updating subtasks in your task plan (e.g., segmenting exploration by batches, categories, or ranges like a-d) explicitly documents the exact paths and progress already completed, maintaining clear self-awareness of your exploration trajectory and preventing backtracking.\n 3. **Single-Pass Trust & Anti-Oscillation (No Excessive Re-verification)**: Trust your recorded observations as the ground truth. Once a section has been traversed and logged, consider it permanently verified. Never scroll back and forth to \"double-check\" or execute redundant full-screen re-scans, unless explicit verification feedback points to a specific missing item or you are actively resolving a rejected state.\n- **Short-Term Memory**: short_term_memory is a supplement to the task plan. If needed, use the tag in your raw thinking text to record it. Short-term memory will be passed to the agent in the next turn, you can describe any trivial details, it should at least include expectations for the future, the current short-term strategy, and verified facts. Example:\n\nExpectation: Duplicate targets were not found in e-h, now start looking for items starting with i-l. (Letters after m are still unverified); Operation strategy: Search in batches by starting letter; Verified facts: a-d, e-h have no matches.\n\n\n# SYSTEM WORKFLOW: SAFETY NET & AUTO-RECOVERY\nYour decisions are processed through downstream pipeline guardians to guarantee accuracy:\n- **Pre-Execution Safety Net**: Automatically checks that your target is visible and safe to interact with prior to action execution.\n- **Failure Analyzer**: Activates if an action fails, running a local step-by-step recovery sequence to self-heal (e.g., closing keyboards or handling wait states).\nThese operations are recorded chronologically in your '--- Execution History ---' to help you trace how the current screen state was reached.\n\n# OBJECTIVE & PLAN\nInitial Goal: {{ initial_goal }}\n\n## Current Plan & Execution History\n{{ plan_and_history }}\n\n# TOOL USE PROTOCOL: PRE-DECISION EXPLORATION VS. TURN-ENDING ACTIONS\nYou operate within a multi-step framework that splits tools into two distinct classes. You must understand this flow to gather details safely before committing to physical actions:\n1. **Pre-Decision Exploratory Tools**:\n - *What they are*: Helper/Subagent tools (such as `ask_explorer`, `ask_diagnoser`, `video_analyzer`), ADB command tools (`run_adb_command`, `manage_task`), and memory note tools (`read_note`, `list_notes`, `save_note`, `update_note`, `append_note`).\n - *How they work*: These tools gather details or update/read memory. When you invoke these tools, the framework will immediately execute them and return the results, allowing you to continue thinking and make your final decision.\n2. **Turn-Ending Actions**:\n - *What they are*: Physical device actions (`click`, `input_text`, `swipe`, `press_key`, `manage_app`, `wait_for_delay`, `long_press`) and flow-ending actions (`reply_to_checker`).\n - *How they work*: These tools interact with the physical device or alter the operational flow. Executing any of these tools terminates your turn.\n\n# THE THREE-PHASE COGNITIVE PROTOCOL\nIn every turn, proceed through these three phases in sequence.\n\n## Phase 1: Observation & Diagnosis\n1. **Visual Observation & State Verification**: You must first examine the screenshot and describe the visual screen state in your raw thinking before evaluating actions.\n - *Visual Description*: Describe the current visual layout, key visible UI components, active screen state, and any unexpected popups or overlays shown on the screenshot.\n - *State Verification*: Compare the observed visual state against the intended outcome of the previous action to confirm whether the transition succeeded.\n - *Side-Effect & Anomaly Check*: Verify whether any unintended screen changes, unwanted dialogs, or unexpected UI state alterations occurred.\n - *Drift/Regression Audit*: Has the UI drifted to an incorrect screen, or Regressed back to a previous screen (e.g. from an accidental browser back-navigation)?\n - *Stuck/Transition state*: Is there a loading progress indicator (active transition) or is it a frozen screen (static stuck)?\n - *Missed Transient States*: If the current state jumped ahead and you suspect you missed an important initial or intermediate screen (e.g., a number or flash message that was overwritten by your last action), consider using the tool to look back at the recent history.\n2. **Path Decision (Recover vs. Adapt)**: Encourage flexible problem-solving. First, evaluate if the current path can still be recovered or advanced (e.g., dismissing a blocker, retrying, or restoring the view). If the current approach is blocked or unreliable, do not get stuckโ€”flexibly pivot to alternative solutions or different UI paths to achieve the same goal.\n3. **Verify & Diagnose (Mandatory Diagnosis Triggers)**: You **MUST** use the diagnostic tool (`ask_diagnoser`) during the Pre-Decision Exploration stage (before calling any Turn-Ending Actions) if you detect any of the following symptoms:\n - **Action Loop or State Stagnation**: You are about to execute the exact same physical action or click the same target for a 3rd consecutive time, or you find yourself alternating between the same actions (e.g., swiping up then down repeatedly) in a cyclical loop. This also applies if the current screen and 'Visible UI Elements' list are virtually identical to the previous step, or oscillating between two identical states, even though you executed a physical action. (Exception: Continuous swiping in one direction during large list/text exploration to view full content or reveal new items is acceptable and exempt, provided new content is being browsed or recorded into notes according to your exploration plan).\n - **Validation/Checker Rejection**: The verification agent (Checker) rejected your subgoal completion (i.e., you are in troubleshooter mode and this is a retry).\n - **Unexpected UI Drift or App Crash**: An unexpected OS popup, permissions dialog, system crash box, or redirect to home screen occurred.\n - **Visual Discrepancies/Uncertainty**: You see a loading progress spinner for more than 2 consecutive steps, or you are unsure why an action failed to trigger the expected screen change.\n Do NOT execute any Turn-Ending Actions in the same turn as a mandatory diagnosis. Call `ask_diagnoser` FIRST, analyze its diagnostic reply, and then perform the corrected action in the next turn.\n\n## Phase 2: Planning & Memory Preservation\n1. **Autonomous Goal Management**: You have complete autonomy over subgoals under the top-level milestones. You MUST actively update the status of the plan and proactively add new subgoals as you progress. Use `[/]` to denote the currently active milestone and subgoal. Only mark a step as completed (`[x]`) AFTER visual confirmation. Proactively maintain the task plan by dynamically adding new subgoals or adjusting existing ones based on your strategy. Keep the plan as a pure Markdown checklist in `task_plan`.\n *Example format (e.g., overall goal is to process items in a large gallery):*\n - [x] Open the gallery on the device\n - [x] Open the app list and tap on gallery\n - [/] Systematic Single-Pass traversal from top to bottom\n - [x] Scan section 1 (items 1-10), log unique identifiers into \"Gallery Scan\" note\n - [x] Scroll using item 10 as anchor, scan section 2 (items 11-20), process duplicates against note\n - [/] Scroll using item 20 as anchor, scan section 3 (items 21 to end)\n - [ ] Finalize task: all items processed in single pass, confirm note summary\n **CRITICAL: When all top-level tasks are marked as completed (`[x]`) using the note tool, the system will automatically judge the task as finished. At this point, you do not need to output a Turn-Ending Action.**\n2. **Preserve Memory & Pre-Action Snapshotting**: You MUST perform the following risk assessment before any UI interaction: \"Will this action irreversibly alter the screen and erase potentially useful information?\" If this risk exists, even if you are not yet certain if you will need these details, you should immediately save the extracted critical values to the note tools to ensure safe downstream access. When writing notes, you must clearly annotate the context in which the information was obtained. Saving more notes will not cause a burden; rather, it prevents risks caused by forgetting. Pay attention to verification when writing notes, but once a final answer is determined, it should be trusted unless there is clear evidence that the existing note is incorrect or outdated. Please use clear key-value structures to organize the data and detail the background. If you have already missed certain critical information and cannot re-view it by going back or other operations, you can use the `video_analyzer` to look back at historical information.\n3. **Background Tasks & Instrumentation (Optional)**: If the current subgoal requires measuring performance, logging power consumption, or recording traces, you may use ADB tools to start these background processes during the Pre-Decision Exploration stage, prior to physical execution. You have full control over all ADB tasks. If a background task has finished, you must understand the original purpose of starting that task, and record the task's results or relevant extracted data in your notes as needed.\n\n## Phase 3: Target Selection & Physical Execution\n1. **Target Resolution & Index Safety**: The indexed list coordinates are generally accurate, but sometimes they include elements that are not fully visible or are obscured on the screen. You must first analyze the current screenshot to determine which indexes are actually interactable, and prioritize using element indexes (e.g., `target=3`) over raw coordinates or cached state whenever possible. If the target is not present in the indexed list or the bounds are incorrect, use `ask_explorer` first during the Pre-Decision Exploration stage. Remember that swiping, scrolling, or page updates shift indexes and coordinates; always retrieve the fresh index from the latest screen data.\n2. **Action Selection & Mandatory Cognitive Protocol**: Before calling any Turn-Ending Action (`click`, `swipe`, `input_text`, `long_press`, `press_key`, `manage_app`, or `wait_for_delay`), you MUST write a single, fluid narrative paragraph of text as your raw thinking. In this paragraph, articulate why this specific tool/target is the safest, most accurate next step, and define the exact immediate screen change you expect as proof of success.", + "troubleshooter_template": "# ROLE\nYou are the Expert Android Device Troubleshooter. A previous attempt to complete the task was rejected by the verification agent (Checker). Your priority is diagnosis and execution accuracy to resolve the rejection.\n\n# OPERATIONAL PRINCIPLES\n- **Action Scope & Interaction Precision**: Always select gestures whose scope of effect matches the boundaries of the target element. Avoid shortcuts or batch actions from popups or menus if they risk affecting non-target UI components on the screen. For any task requiring continuous range interactions, sliding, or boundary selection, prefer using targeted swipe/drag gestures spanning from the element's start bounds [xmin, ymin] to end bounds [xmax, ymax] to ensure the action is strictly confined to the target element. For drag-and-drop, list reordering, or sliding sliders/SeekBars, set duration >= 1000 (e.g. 1500) and always drag slightly PAST the target position to overcome touch slop and reliably trigger the drop/snap/value update. When setting a slider to Maximum (100%) or Minimum (0%), always swipe fully to the extreme boundary to guarantee reaching the absolute limits.\n- **Vision-First Priority**: The screenshot is the primary source of truth for the real device state. Always visually inspect the screen state and layout on the screenshot before referencing the element list. Use the 'Visible UI Elements' list strictly as an auxiliary lookup table to retrieve target element indices and coordinates.\n- **Tool Calling Iron Rule**: Do NOT submit a Turn-Ending Action at the same time (in the same tool-call list) as a Pre-Decision Exploratory Tool. You must first gather information (e.g., finding coordinates or diagnosing a state), manage ADB background tasks, or update memory notes. You can use pre-decision exploratory tools to complete this task. You will be prompted again with the tool's result, at which point you can output your final physical/turn-ending action.\n- **Loading & Transitions**: When entering page transitions (e.g., app launching, web page loading, or progress indicators), use `wait_for_delay` to let the screen settle before executing target actions. Some transition states are difficult to discern, such as a thin progress line at the top of the browser screen during webpage loads. Before questioning whether an error has occurred, you can first adopt a waiting strategy to see how the device state evolves. If the state remains ambiguous after waiting, certain UI components might be malfunctioning or failing to load; in this case, consider retrying the action or pivoting to an alternative path.\n- **Memory Preservation & Paginated Exploration**: Using note tools is entirely zero-cost. Actively use this tool to maintain goals and logs. The note tool is the only reliable way to persist memory. Therefore, when you find information on the screen that might be relevant to the task, even if you are not certain, you should first save the information currently visible on the screen before exploring further. If you have already recorded some information, you can use `append_note` to append the new information to the existing content. The note tool can also help you organize lists or information across multiple screens; you can consolidate all information into notes for unified processing.\n- **Large List/Text/Complex UI Exploration Strategy (Zero Omission, Zero Duplication, Single-Pass Rule)**:\n When traversing large lists, long text bodies, or continuous content streams, adhere strictly to the **Single-Pass Exploration Standard**:\n 1. **Unidirectional Anchor-Based Traversal (Zero Omission)**: Maintain a consistent, unidirectional scan (e.g., scrolling steadily from top to bottom). Before each scroll, use the bottom-most visible item as a visual anchor. Ensure your scrolling distance allows this anchor to remain visible near the top of the subsequent screen, providing seamless visual overlap so no items are jumped over or omitted. Conclude exploration when reaching the definite boundary (e.g., reaching the bottom of the list or when further scrolling reveals no new content).\n 2. **Memory & Plan-Driven Deduplication (Zero Duplication)**: Systematically log extracted item names, identifiers, or processed statuses into notes or memory as you progress. Upon every screen update, cross-reference visible items with your recorded data and strictly process only *newly surfaced, unrecorded* items to avoid duplicate actions on overlapping elements. Furthermore, actively maintaining and updating subtasks in your task plan (e.g., segmenting exploration by batches, categories, or ranges like a-d) explicitly documents the exact paths and progress already completed, maintaining clear self-awareness of your exploration trajectory and preventing backtracking.\n 3. **Single-Pass Trust & Anti-Oscillation (No Excessive Re-verification)**: Trust your recorded observations as the ground truth. Once a section has been traversed and logged, consider it permanently verified. Never scroll back and forth to \"double-check\" or execute redundant full-screen re-scans, unless explicit verification feedback points to a specific missing item or you are actively resolving a rejected state.\n- **Short-Term Memory**: short_term_memory is a supplement to the task plan. If needed, use the tag in your raw thinking text to record it. Short-term memory will be passed to the agent in the next turn, you can describe any trivial details, it should at least include expectations for the future, the current short-term strategy, and verified facts. Example:\n\nExpectation: Duplicate targets were not found in e-h, now start looking for items starting with i-l. (Letters after m are still unverified); Operation strategy: Search in batches by starting letter; Verified facts: a-d, e-h have no matches.\n\n\n# SYSTEM WORKFLOW: SAFETY NET & AUTO-RECOVERY\nYour decisions are processed through downstream pipeline guardians to guarantee accuracy:\n- **Pre-Execution Safety Net**: Automatically checks that your target is visible and safe to interact with prior to action execution.\n- **Failure Analyzer**: Activates if an action fails, running a local step-by-step recovery sequence to self-heal (e.g., closing keyboards or handling wait states).\nThese operations are recorded chronologically in your '--- Execution History ---' to help you trace how the current screen state was reached.\n\n# OBJECTIVE & PLAN\nInitial Goal: {{ initial_goal }}\n\n## Current Plan & Execution History\n{{ plan_and_history }}\n\n# TOOL USE PROTOCOL: PRE-DECISION EXPLORATION VS. TURN-ENDING ACTIONS\nYou operate within a multi-step framework that splits tools into two distinct classes. You must understand this flow to gather details safely before committing to physical actions:\n1. **Pre-Decision Exploratory Tools**:\n - *What they are*: Helper/Subagent tools (such as `ask_explorer`, `ask_diagnoser`, `video_analyzer`), ADB command tools (`run_adb_command`, `manage_task`), and memory note tools (`read_note`, `list_notes`, `save_note`, `update_note`, `append_note`).\n - *How they work*: These tools gather details or update/read memory. When you invoke these tools, the framework will immediately execute them and return the results, allowing you to continue thinking and make your final decision.\n2. **Turn-Ending Actions**:\n - *What they are*: Physical device actions (`click`, `input_text`, `swipe`, `press_key`, `manage_app`, `wait_for_delay`, `long_press`) and flow-ending actions (`reply_to_checker`).\n - *How they work*: These tools interact with the physical device or alter the operational flow. Executing any of these tools terminates your turn.\n\n# THE THREE-PHASE COGNITIVE PROTOCOL\nIn every turn, proceed through these three phases in sequence.\n\n## Phase 1: Observation & Diagnosis\n1. **Analyze Transition & Rejection**: Compare the current screen and the 'Visible UI Elements' list with the last step's action context, and read the Checker's feedback to evaluate the discrepancy and understand the rejection reasons. If there is short-term memory, that will also be highly valuable information.\n - *Transition & Rejection Analysis*: Identify the exact changes between screen states and note potential unexpected events (app crashes, autoplay transitions, redirects, or UI state drift).\n - *Drift/Regression Audit*: Spot if the current screen has drifted to an unintended view, or unexpectedly Regressed back to an earlier page we already left (e.g., due to an accidental browser back-navigation).\n - *Stuck/Transition state*: Is there a loading progress indicator (active transition) or is it a frozen/stuck screen?\n - *Missed Transient States*: If the current state jumped ahead and you suspect you missed an important initial or intermediate screen (e.g., a number or flash message that was overwritten by your last action), consider using the tool to look back at the recent history.\n - *Challenge Checker*: If you visually confirm the subgoal's target state is *already achieved*, you **MUST** invoke the `reply_to_checker` tool to state your observation.\n2. **Path Decision (Recover vs. Adapt)**: Evaluate why the previous attempt failed and explore multiple problem-solving paths. First, determine if the original path can still be salvaged (e.g., resubmitting input correctly, bypassing a blocker, or re-navigating). If the original method is unreliable, do not stubbornly repeat itโ€”flexibly pivot to an alternative approach or a different UI path to satisfy the Checker. Save discoveries and lessons in notes.\n3. **Verify & Diagnose (Mandatory Diagnosis Triggers)**: You **MUST** use the diagnostic tool (`ask_diagnoser`) during the Pre-Decision Exploration stage (before calling any Turn-Ending Actions) if you detect any of the following symptoms:\n - **Action Loop or State Stagnation**: You are about to execute the exact same physical action or click the same target for a 3rd consecutive time, or you find yourself alternating between the same actions (e.g., swiping up then down repeatedly) in a cyclical loop. This also applies if the current screen and 'Visible UI Elements' list are virtually identical to the previous step, or oscillating between two identical states, even though you executed a physical action. (Exception: Continuous swiping in one direction during large list/text exploration to view full content or reveal new items is acceptable and exempt, provided new content is being browsed or recorded into notes according to your exploration plan).\n - **Validation/Checker Rejection**: The verification agent (Checker) rejected your subgoal completion (i.e., you are in troubleshooter mode and this is a retry).\n - **Unexpected UI Drift or App Crash**: An unexpected OS popup, permissions dialog, system crash box, or redirect to home screen occurred.\n - **Visual Discrepancies/Uncertainty**: You see a loading progress spinner for more than 2 consecutive steps, or you are unsure why an action failed to trigger the expected screen change.\n Do NOT execute any Turn-Ending Actions in the same turn as a mandatory diagnosis. Call `ask_diagnoser` FIRST, analyze its diagnostic reply, and then perform the corrected action in the next turn.\n\n## Phase 2: Planning & Memory Preservation\n1. **Autonomous Goal Management**: You have complete autonomy over subgoals under the top-level milestones. You MUST actively update the status of the plan and proactively add new subgoals as you progress. Use `[/]` to denote the currently active milestone and subgoal. Only mark a step as completed (`[x]`) AFTER visual confirmation. Proactively maintain the task plan by dynamically adding new subgoals or adjusting existing ones based on your strategy. Keep the plan as a pure Markdown checklist in `task_plan`.\n *Example format (e.g., overall goal is to process items in a large gallery):*\n - [x] Open the gallery on the device\n - [x] Open the app list and tap on gallery\n - [/] Systematic Single-Pass traversal from top to bottom\n - [x] Scan section 1 (items 1-10), log unique identifiers into \"Gallery Scan\" note\n - [x] Scroll using item 10 as anchor, scan section 2 (items 11-20), process duplicates against note\n - [/] Scroll using item 20 as anchor, scan section 3 (items 21 to end)\n - [ ] Finalize task: all items processed in single pass, confirm note summary\n **CRITICAL: When all top-level tasks are marked as completed (`[x]`) using the note tool, the system will automatically judge the task as finished. At this point, you do not need to output a Turn-Ending Action.**\n2. **Preserve Memory & Pre-Action Snapshotting**: You MUST perform the following risk assessment before any UI interaction: \"Will this action irreversibly alter the screen and erase potentially useful information?\" If this risk exists, even if you are not yet certain if you will need these details, you should immediately save the extracted critical values to the note tools to ensure safe downstream access. When writing notes, you must clearly annotate the context in which the information was obtained. Saving more notes will not cause a burden; rather, it prevents risks caused by forgetting. Pay attention to verification when writing notes, but once a final answer is determined, it should be trusted unless there is clear evidence that the existing note is incorrect or outdated. Please use clear key-value structures to organize the data and detail the background. If you have already missed certain critical information and cannot re-view it by going back or other operations, you can use the `video_analyzer` to look back at historical information.\n3. **Background Tasks & Instrumentation (Optional)**: If the current subgoal requires measuring performance, logging power consumption, or recording traces, you may use ADB tools to start these background processes during the Pre-Decision Exploration stage, prior to physical execution. You have full control over all ADB tasks. If a background task has finished, you must understand the original purpose of starting that task, and record the task's results or relevant extracted data in your notes as needed.\n\n## Phase 3: Target Selection & Physical Execution\n1. **Target Resolution & Index Safety**: The indexed list coordinates are generally accurate, but sometimes they include elements that are not fully visible or are obscured on the screen. You must first analyze the current screenshot to determine which indexes are actually interactable, and prioritize using element indexes (e.g., `target=3`) over raw coordinates or cached state whenever possible. If the target is not present in the indexed list or the bounds are incorrect, use `ask_explorer` first during the Pre-Decision Exploration stage. Remember that swiping, scrolling, or page updates shift indexes and coordinates; always retrieve the fresh index from the latest screen data.\n2. **Action Selection & Mandatory Cognitive Protocol**: Before calling any Turn-Ending Action (`click`, `swipe`, `input_text`, `long_press`, `press_key`, `manage_app`, or `wait_for_delay`), you MUST write a single, fluid narrative paragraph of text as your raw thinking. In this paragraph, articulate why this specific tool/target is the safest, most accurate next step, and define the exact immediate screen change you expect as proof of success." +} diff --git a/artemis/agents/operator/operator.py b/artemis/agents/operator/operator.py new file mode 100644 index 0000000..33babf8 --- /dev/null +++ b/artemis/agents/operator/operator.py @@ -0,0 +1,1481 @@ +# 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 asyncio +import json +from pathlib import Path +import re +from typing import Annotated, Any, Literal +from uuid import uuid4 + +from langchain_core.messages import SystemMessage, ToolMessage +from langchain_core.tools import BaseTool, tool +from langgraph.types import Command + +from artemis.constants import ( + OPERATOR_MAX_CHAT_ROUNDS, + VALIDATOR_MESSAGES_KEY, +) +from artemis.context import ArtemisContext +from artemis.data_engine.trace import ( + TraceSpan, + trace, + trace_langchain_tool, +) +from artemis.graph.state import State +from artemis.services.llm import get_llm, invoke_llm_with_timeout_message +from artemis.tools.command_tool import ( + _BACKGROUND_TASKS, + _FINISHED_TASKS_LOGS, + _is_output_long, + analyze_task_output_wrapper, +) +from artemis.tools.index import get_tool_by_name +from artemis.tools.tool_wrapper import ( + get_tool_result_content, + invoke_tool_with_injection, +) +from artemis.utils.coordinates import parse_swipe_parameters +from artemis.utils.decorators import wrap_with_callbacks +from artemis.utils.file import create_snapshot +from artemis.utils.logger import get_logger +from artemis.utils.notes import get_note_file_path, get_notes_dir +from artemis.utils.task_tree import ( + build_plan_and_history, + get_active_subgoal_hashes, +) +from artemis.utils.verification import ( + append_verification_chat, + get_verification_chat_path, + get_verification_chat_rounds, + read_verification_chat, +) +from artemis.utils.visualization import format_minimal_list_with_elements + +logger = get_logger(__name__) + +DEFERRING_TOOLS = { + "ask_diagnoser", + "video_analyzer", + "read_note", + "list_notes", + "save_note", + "run_adb_command", + "manage_task", + "analyze_task_output", + "ask_explorer", +} + +from artemis.agents.operator.prompts import ( + PromptBuilder, + PromptComponent, + TemplatePromptComponent, + ObservationPromptComponent, + ScreenshotSimilarityPromptComponent, + CheckerFeedbackPromptComponent, + BackgroundTasksPromptComponent, + ShortTermMemoryPromptComponent, + TaskPlanWarningPromptComponent, + ToolLimitWarningPromptComponent, + InjectedInstructionPromptComponent, +) + + +class OperatorNode: + def __init__( + self, + ctx: ArtemisContext, + tools: list[BaseTool] | None = None, + prompt_components: list[PromptComponent] | None = None, + last_n_detailed: int = 1, + ): + self.ctx = ctx + self.tools = tools or [] + self.prompt_components = prompt_components or [] + self.last_n_detailed = last_n_detailed + self.prompts = {} + prompts_path = Path(__file__).parent / "operator.json" + if prompts_path.exists(): + try: + self.prompts = json.loads(prompts_path.read_text(encoding="utf-8")) + except Exception as e: + logger.error(f"Failed to load operator prompts: {e}") + + def _check_infinite_loop(self, state: State): + subagent_calls = state.subagent_calls or [] + if len(subagent_calls) >= 3: + last_three = subagent_calls[-3:] + if len(set(last_three)) == 1: + logger.error( + f"Infinite loop detected: Sub-agent {last_three[0]} called" + " 3 times consecutively." + ) + raise RuntimeError( + f"Infinite loop detected: Sub-agent {last_three[0]} called" + " 3 times consecutively." + ) + + def _get_history_and_plan(self) -> tuple[list[dict], str]: + steps = [] + if self.ctx.data_engine: + steps = self.ctx.data_engine.get_agent_friendly_steps() + + task_plan = "No task plan yet." + if self.ctx.data_engine: + task_plan_path = get_note_file_path(self.ctx.data_engine.base_dir, "task_plan") + + if task_plan_path.exists(): + try: + task_plan = task_plan_path.read_text(encoding="utf-8") + except Exception as e: + task_plan = f"Error reading task plan: {e}" + + return steps, task_plan + + def _get_verification_chat_rounds(self) -> tuple[int, int]: + if not self.ctx.data_engine: + return 0, 0 + + subgoal_hash = self._get_active_subgoal_hash() + chat_path = get_verification_chat_path(self.ctx.data_engine.base_dir, subgoal_hash) + turns = read_verification_chat(chat_path) + return get_verification_chat_rounds(turns) + + def _has_checker_feedback(self) -> bool: + if not self.ctx.data_engine: + return False + + subgoal_hash = self._get_active_subgoal_hash() + chat_path = get_verification_chat_path(self.ctx.data_engine.base_dir, subgoal_hash) + return chat_path.exists() + + def _get_reply_to_checker_tool(self) -> BaseTool: + @tool + def reply_to_checker( + reasoning: Annotated[ + str, + "Your reasoning or observations regarding the verification feedback.", + ], + ): + """[TERMINAL] Use this tool to reply to the verification feedback with your reasoning or observations. + + This ends your turn. + """ + + if self.ctx.data_engine: + subgoal_hash = self._get_active_subgoal_hash() + chat_path = get_verification_chat_path(self.ctx.data_engine.base_dir, subgoal_hash) + + turns = read_verification_chat(chat_path) + max_op, max_chk = get_verification_chat_rounds(turns) + + round_num = max(max_op + 1, max_chk) + if round_num == 0: + round_num = 1 + + if append_verification_chat(chat_path, "operator", reasoning, round_num): + return "Successfully recorded your reply to the Checker." + else: + return "Failed to save reply." + else: + return "Error: DataEngine not available." + + return reply_to_checker + + async def _build_prompt( + self, + state: State, + latest_screenshot_b64: str, + fused_xml: list[dict], + minimal_list: str, + current_step_num: int, + steps: list[dict], + task_plan: str, + active_background_tasks: list[dict] = None, + newly_finished_tasks: list[dict] = None, + ) -> list: + plan_and_history = self._build_plan_and_history(steps, task_plan) + + builder = PromptBuilder() + + has_feedback = self._has_checker_feedback() + + components = self.prompt_components + if not components: + if has_feedback: + components = [ + ( + TemplatePromptComponent(), + {"template_name": "troubleshooter_template"}, + ), + (ObservationPromptComponent(), {}), + (ScreenshotSimilarityPromptComponent(), {}), + (InjectedInstructionPromptComponent(), {}), + (ShortTermMemoryPromptComponent(), {}), + (BackgroundTasksPromptComponent(), {}), + (CheckerFeedbackPromptComponent(), {}), + (TaskPlanWarningPromptComponent(), {}), + (ToolLimitWarningPromptComponent(), {}), + ] + else: + components = [ + ( + TemplatePromptComponent(), + {"template_name": "main_template"}, + ), + (ObservationPromptComponent(), {}), + (ScreenshotSimilarityPromptComponent(), {}), + (InjectedInstructionPromptComponent(), {}), + (ShortTermMemoryPromptComponent(), {}), + (BackgroundTasksPromptComponent(), {}), + (TaskPlanWarningPromptComponent(), {}), + (ToolLimitWarningPromptComponent(), {}), + ] + else: + components = [(c, {}) for c in components] + [(ToolLimitWarningPromptComponent(), {})] + + # Active and finished tasks are now passed as arguments + active_tasks = active_background_tasks or [] + newly_finished = newly_finished_tasks or [] + + for component, extra_kwargs in components: + kwargs_to_pass = { + "prompts": self.prompts, + "plan_and_history": plan_and_history, + "latest_screenshot_b64": latest_screenshot_b64, + "minimal_list": minimal_list, + "current_step_num": current_step_num, + "active_background_tasks": active_tasks, + "newly_finished_tasks": newly_finished, + "steps": steps, + } + kwargs_to_pass.update(extra_kwargs) + await component(builder, state, self.ctx, **kwargs_to_pass) + + return builder.build() + + async def _invoke_llm_loop( + self, + base_llm, + current_messages: list, + traced_tools: list, + new_subagent_calls: list, + state: State, + ) -> tuple[list[dict] | None, bool, str | None, str | None, bool]: + max_iterations = 20 + action_result = None + requested_argue = False + raw_thoughts = [] + native_thoughts = [] + tool_limit_exceeded = False + + async def read_stream_and_accumulate(target_llm, msgs): + try: + full_response = None + has_chunks = False + async for chunk in target_llm.astream(msgs): + has_chunks = True + if full_response is None: + full_response = chunk + else: + full_response += chunk + if has_chunks and full_response is not None: + return full_response + except Exception as e: + logger.debug(f"astream failed or not supported: {e}. Falling back to ainvoke.") + + return await target_llm.ainvoke(msgs) + + for iteration in range(max_iterations): + if iteration == max_iterations - 1: + logger.warning( + "Operator has reached the maximum number of tool calls. Adding reminder." + ) + current_messages.append( + SystemMessage( + content=( + "You have reached the maximum number of tool calls" + " for this turn. You must make a final action, or" + " save your current research results via notes to" + " continue later." + ) + ) + ) + + bound_llm = base_llm.bind_tools(tools=traced_tools) + response = await invoke_llm_with_timeout_message( + read_stream_and_accumulate(bound_llm, current_messages) + ) + + if hasattr(response, "response_metadata") and response.response_metadata: + usage = ( + response.response_metadata.get("usage_metadata") + or response.response_metadata.get("token_usage") + or {} + ) + prompt_tokens = usage.get("prompt_token_count") or usage.get("prompt_tokens") + cached_tokens = usage.get("cached_content_token_count") or usage.get( + "cached_tokens", 0 + ) + completion_tokens = usage.get("candidates_token_count") or usage.get( + "completion_tokens" + ) + if prompt_tokens is not None: + logger.info( + f"LLM usage: prompt_tokens={prompt_tokens}" + f" (cached={cached_tokens})," + f" completion_tokens={completion_tokens}" + ) + + if response.content: + if isinstance(response.content, str): + raw_thoughts.append(response.content) + elif isinstance(response.content, list): + for item in response.content: + if isinstance(item, dict): + if item.get("type") == "text": + raw_thoughts.append(item["text"]) + elif item.get("type") == "thinking": + native_thoughts.append(item["thinking"]) + + additional_kwargs = getattr(response, "additional_kwargs", None) + if ( + not response.tool_calls + and isinstance(additional_kwargs, dict) + and not hasattr(additional_kwargs, "assert_called_with") # Filter out mocks + and additional_kwargs.get("function_call") + ): + fc = additional_kwargs["function_call"] + try: + arguments_str = fc.get("arguments", "{}") + if isinstance(arguments_str, str): + arguments = json.loads(arguments_str) + else: + arguments = arguments_str + + tool_call = { + "name": fc["name"], + "args": arguments, + "id": f"call_{uuid4().hex[:8]}", + "type": "tool_call", + } + response.tool_calls = [tool_call] + logger.info( + "Fallback: Manually parsed function_call from" + f" additional_kwargs: {tool_call}" + ) + except Exception as e: + logger.error(f"Failed to parse function_call fallback: {e}") + + if not response.tool_calls: + logger.warning("LLM stopped without calling any tool. Encouraging action.") + break + + current_messages.append(response) + + action_tool_names = [ + "click", + "input_text", + "swipe", + "press_key", + "manage_app", + "wait_for_delay", + "long_press", + ] + + def normalize_name(name: str) -> str: + return name.split(":")[-1] if ":" in name else name + + action_calls = [ + tc for tc in response.tool_calls if normalize_name(tc["name"]) in action_tool_names + ] + reply_to_checker_calls = [ + tc for tc in response.tool_calls if normalize_name(tc["name"]) == "reply_to_checker" + ] + other_calls = [ + tc + for tc in response.tool_calls + if normalize_name(tc["name"]) not in action_tool_names + ["reply_to_checker"] + ] + + tool_outputs = [] + validation_errors = False + requested_argue = False + + other_tool_failed = False + other_tool_failure_msg = "" + + has_terminal_calls = bool(action_calls or reply_to_checker_calls) + + if other_calls: + results = [] + for tc in other_calls: + tool_name = tc["name"] + logger.info(f"Operator requested tool: {tool_name}") + if tool_name in ["ask_diagnoser", "video_analyzer"]: + new_subagent_calls.append(tool_name) + + if ":" in tool_name: + tool_to_run = get_tool_by_name(tool_name, traced_tools) + else: + tool_to_run = next( + (t for t in traced_tools if t.name == tool_name), + None, + ) + + if tool_to_run: + args = dict(tc["args"]) + span = TraceSpan(name=tool_name, trace_type="tool", ctx=self.ctx) + span.payload = {"args": args} + + with span: + try: + state.operator_raw_thinking = ( + "\n".join(raw_thoughts) if raw_thoughts else None + ) + state.operator_native_thinking = ( + "\n".join(native_thoughts) if native_thoughts else None + ) + + result_obj = await invoke_tool_with_injection( + tool=tool_to_run, + args=args, + tool_call_id=tc["id"], + state=state, + ) + + content = get_tool_result_content(result_obj) + status = "success" + + if isinstance(result_obj, Command): + updates = result_obj.update + if VALIDATOR_MESSAGES_KEY in updates: + msgs = updates[VALIDATOR_MESSAGES_KEY] + if msgs and hasattr(msgs[0], "status"): + status = msgs[0].status + + is_err = status == "error" + if not is_err: + if isinstance(content, str): + if content.startswith( + "Failed" + ) or content.lower().startswith("error"): + is_err = True + elif isinstance(content, list): + for block in content: + if ( + isinstance(block, dict) + and block.get("type") == "text" + ): + text = block.get("text", "") + if text.startswith( + "Failed" + ) or text.lower().startswith("error"): + is_err = True + break + + if is_err: + other_tool_failed = True + other_tool_failure_msg = str(content) + status = "error" + + span.status = "failed" + span.error = str(content) + else: + span.result = content + + results.append((content, status)) + except Exception as e: + other_tool_failed = True + other_tool_failure_msg = str(e) + + span.status = "failed" + span.error = str(e) + + results.append( + ( + f"Error running tool {tool_name}: {e}", + "error", + ) + ) + else: + other_tool_failed = True + other_tool_failure_msg = f"Tool {tool_name} not supported" + results.append((f"Error: Tool {tool_name} not supported", "error")) + + for tc, (content, status) in zip(other_calls, results): + tool_outputs.append( + ToolMessage( + tool_call_id=tc["id"], + content=content, + status=status, + ) + ) + + if other_tool_failed: + if action_calls: + logger.warning( + "Scratchpad/helper tool call failed. Deferring" + f" screen actions: {other_tool_failure_msg}" + ) + for tc in action_calls: + tool_outputs.append( + ToolMessage( + tool_call_id=tc["id"], + content=( + "This screen action was rejected" + " because an accompanying tool call" + " failed. Error:" + f" {other_tool_failure_msg}. Currently," + " no screen actions have been" + " executed. Please correct the error" + " and try again." + ), + status="error", + ) + ) + action_calls = [] + + if reply_to_checker_calls: + logger.warning( + "Scratchpad/helper tool call failed. Deferring" + f" reply to checker: {other_tool_failure_msg}" + ) + for tc in reply_to_checker_calls: + tool_outputs.append( + ToolMessage( + tool_call_id=tc["id"], + content=( + "Your reply to the checker was" + " rejected because an accompanying" + " tool call failed. Error:" + f" {other_tool_failure_msg}. Currently," + " no response has been sent to the" + " checker. Please correct the error" + " and try again." + ), + status="error", + ) + ) + reply_to_checker_calls = [] + + elif has_terminal_calls: + has_deferring_calls = any( + normalize_name(tc["name"]) in DEFERRING_TOOLS for tc in other_calls + ) + if has_deferring_calls: + logger.info( + "Operator mixed pre-decision tools" + " (memory/exploration) and terminal action calls." + " Deferring terminal actions." + ) + + for tc in action_calls: + tool_outputs.append( + ToolMessage( + tool_call_id=tc["id"], + content=( + "Your pre-decision tools have been" + " successfully processed. However, the" + " screen action you attempted to" + " execute was rejected because you are" + " not allowed to simultaneously" + " read/write memory and execute" + " physical actions in the same turn." + " Currently, no screen actions have" + " been executed. Please review your" + " updated context and re-output your" + " intended screen action." + ), + status="success", + ) + ) + + for tc in reply_to_checker_calls: + tool_outputs.append( + ToolMessage( + tool_call_id=tc["id"], + content=( + "Your pre-decision tools have been" + " successfully processed. However, your" + " response to the checker was rejected" + " because you are not allowed to" + " simultaneously read/write memory and" + " reply to the checker in the same" + " turn. Currently, no response has been" + " sent to the checker. Please review" + " your updated context and re-output" + " your reply." + ), + status="success", + ) + ) + + action_calls = [] + reply_to_checker_calls = [] + + if action_calls: + logger.info( + f"Operator requested {len(action_calls)} action(s)." + " Translating and processing all." + ) + actions_list = [] + for tc in action_calls: + translated_actions, error = self._translate_and_validate_tool(tc, state) + if error: + validation_errors = True + tool_outputs.append( + ToolMessage( + tool_call_id=tc["id"], + content=error, + status="error", + ) + ) + else: + actions_list.extend(translated_actions) + tool_outputs.append( + ToolMessage( + tool_call_id=tc["id"], + content="Action Recorded", + status="success", + ) + ) + + if not validation_errors: + action_result = actions_list + + if reply_to_checker_calls: + logger.info("Operator replied to checker.") + tc = reply_to_checker_calls[0] + requested_argue = True + + # Execute it immediately + result = await self._run_other_tool( + tc, + traced_tools, + new_subagent_calls, + state, + raw_thoughts, + native_thoughts, + ) + tool_outputs.append( + ToolMessage( + tool_call_id=tc["id"], + content=result, + status="success" if not result.startswith("Error") else "error", + ) + ) + + for tm in tool_outputs: + current_messages.append(tm) + + if (action_calls and not validation_errors) or requested_argue: + break + + if validation_errors: + logger.warning("Some action calls failed validation. Feeding back to LLM.") + continue + else: + tool_limit_exceeded = True + + raw_thinking = "\n".join(raw_thoughts) if raw_thoughts else None + native_thinking = "\n".join(native_thoughts) if native_thoughts else None + return ( + action_result, + requested_argue, + raw_thinking, + native_thinking, + tool_limit_exceeded, + ) + + async def _run_other_tool( + self, + tc, + traced_tools, + new_subagent_calls, + state: State, + raw_thoughts: list, + native_thoughts: list, + ) -> str: + tool_name = tc["name"] + logger.info(f"Operator requested tool: {tool_name}") + if tool_name in ["ask_diagnoser", "video_analyzer"]: + new_subagent_calls.append(tool_name) + if ":" in tool_name: + tool_to_run = get_tool_by_name(tool_name, traced_tools) + else: + tool_to_run = next((t for t in traced_tools if t.name == tool_name), None) + if tool_to_run: + try: + args = dict(tc["args"]) + + state.operator_raw_thinking = "\n".join(raw_thoughts) if raw_thoughts else None + state.operator_native_thinking = ( + "\n".join(native_thoughts) if native_thoughts else None + ) + + result_obj = await invoke_tool_with_injection( + tool=tool_to_run, + args=args, + tool_call_id=tc["id"], + state=state, + ) + return get_tool_result_content(result_obj) + except Exception as e: + return f"Error running tool {tool_name}: {e}" + else: + return f"Error: Tool {tool_name} not supported" + + @wrap_with_callbacks( + before=lambda: logger.info("Starting Operator Agent..."), + on_success=lambda _: logger.success("Operator Agent"), + on_failure=lambda _: logger.error("Operator Agent"), + ) + @trace(type="agent", name="operator") + async def __call__(self, state: State): + # 1. Check for infinite loops + self._check_infinite_loop(state) + + # Create snapshot for optimistic execution + if self.ctx.data_engine: + notes_dir = get_notes_dir(self.ctx.data_engine.base_dir) + if notes_dir.exists(): + snapshot_dir = notes_dir.with_name("notes_snapshot") + await asyncio.to_thread(create_snapshot, notes_dir, snapshot_dir) + self.ctx.task_plan_snapshot = snapshot_dir + logger.info(f"Created snapshot of notes at {snapshot_dir}") + + # 2. Get perception data from state (populated by Perception node) + operator_raw_data = getattr(state, "operator_raw_data", {}) or {} + latest_screenshot_b64 = operator_raw_data.get("screenshot_b64") + xml_hierarchy = operator_raw_data.get("xml_hierarchy") + operator_raw_data.get("ocr_results") + fused_xml = state.latest_ui_hierarchy if hasattr(state, "latest_ui_hierarchy") else [] + state.latest_screenshot if hasattr(state, "latest_screenshot") else None + + if not latest_screenshot_b64 or xml_hierarchy is None: + logger.error( + "Operator requires perception data in state (populated by Perception node)" + ) + raise ValueError("Operator requires perception data in state") + + # 3. Get history and plan + steps, task_plan = self._get_history_and_plan() + + # 4. Format minimal list + + width = 1080 + height = 2400 + w_raw = operator_raw_data.get("width") + h_raw = operator_raw_data.get("height") + if isinstance(w_raw, int) and isinstance(h_raw, int): + width = w_raw + height = h_raw + else: + try: + if self.ctx and hasattr(self.ctx, "device") and self.ctx.device: + w = getattr(self.ctx.device, "device_width", 1080) + h = getattr(self.ctx.device, "device_height", 2400) + if isinstance(w, int): + width = w + if isinstance(h, int): + height = h + except Exception: + pass + minimal_list, elements, labels = format_minimal_list_with_elements(fused_xml, width, height) + + # Save transient indexed_points and indexed_elements on the state instance + state.indexed_points = [el["center"] for el in elements] + state.indexed_elements = elements + + current_step_num = len(steps) + 1 + + # 5. Prepare Tools + click_tool = self._get_click_tool() + input_text_tool = self._get_input_text_tool() + swipe_tool = self._get_swipe_tool() + press_key_tool = self._get_press_key_tool() + manage_app_tool = self._get_manage_app_tool() + wait_for_delay_tool = self._get_wait_for_delay_tool() + long_press_tool = self._get_long_press_tool() + + all_tools = [ + click_tool, + input_text_tool, + swipe_tool, + press_key_tool, + manage_app_tool, + wait_for_delay_tool, + long_press_tool, + ] + self.tools + + if self._has_checker_feedback(): + max_op, max_chk = self._get_verification_chat_rounds() + max_rounds = ( + getattr( + getattr(self.ctx, "execution_setup", None), + "checker_max_chat_rounds", + OPERATOR_MAX_CHAT_ROUNDS, + ) + if self.ctx and getattr(self.ctx, "execution_setup", None) + else OPERATOR_MAX_CHAT_ROUNDS + ) + if max_op < max_rounds: + reply_to_checker_tool = self._get_reply_to_checker_tool() + all_tools.append(reply_to_checker_tool) + else: + logger.info( + f"Maximum chat rounds ({max_rounds}) reached. Removing reply_to_checker tool." + ) + + # 5. Evaluate dynamic tools and prepare background tasks + has_long_output = False + for tinfo in _FINISHED_TASKS_LOGS.values(): + if _is_output_long(tinfo.get("output", "")): + has_long_output = True + break + if not has_long_output: + for t in _BACKGROUND_TASKS.values(): + if _is_output_long("".join(t.stdout_log)): + has_long_output = True + break + + if has_long_output and not any(t.name == "analyze_task_output" for t in all_tools): + analyze_tool_fn = analyze_task_output_wrapper.tool_fn_getter(self.ctx) + all_tools.append(analyze_tool_fn) + + traced_tools = [trace_langchain_tool(t, self.ctx) for t in all_tools] + + active_background_tasks = [] + for tid, t in _BACKGROUND_TASKS.items(): + active_background_tasks.append( + { + "task_id": tid, + "command": t.command, + "cwd": t.cwd, + "terminal_id": t.terminal_id, + "output_line_count": len(t.stdout_log), + } + ) + + newly_finished_tasks = [] + for tid, tinfo in _FINISHED_TASKS_LOGS.items(): + if not tinfo.get("notified", False): + newly_finished_tasks.append( + { + "task_id": tid, + "command": tinfo.get("command", ""), + "status": tinfo.get("status", "completed"), + "output_text": tinfo.get("output", ""), + } + ) + tinfo["notified"] = True # Mark as read + + # 6. Prepare LLM + llm = get_llm(ctx=self.ctx, name="operator") + + # 7. Build Prompt + messages = await self._build_prompt( + state, + latest_screenshot_b64, + fused_xml, + minimal_list, + current_step_num, + steps, + task_plan, + active_background_tasks=active_background_tasks, + newly_finished_tasks=newly_finished_tasks, + ) + + # 8. Invoke LLM Loop + new_subagent_calls = [] + ( + action_result, + requested_argue, + raw_thinking, + native_thinking, + tool_limit_exceeded, + ) = await self._invoke_llm_loop( + base_llm=llm, + current_messages=messages, + traced_tools=traced_tools, + new_subagent_calls=new_subagent_calls, + state=state, + ) + + # 9. Update State + structured_decisions = json.dumps(action_result) if action_result else None + + short_term_memory = None + combined_thinking = "" + if raw_thinking: + combined_thinking += raw_thinking + "\n" + if native_thinking: + combined_thinking += native_thinking + + if combined_thinking: + stm_match = re.search( + r"(.*?)", + combined_thinking, + re.DOTALL, + ) + if stm_match: + short_term_memory = stm_match.group(1).strip() + + return await state.asanitize_update( + ctx=self.ctx, + update={ + "structured_decisions": structured_decisions, + "operator_raw_thinking": raw_thinking, + "operator_native_thinking": native_thinking, + "short_term_memory": short_term_memory, + "indexed_points": state.indexed_points, + "complete_subgoals_by_ids": [], + "current_step_id": state.current_step_id, + "subagent_calls": ((state.subagent_calls or []) + new_subagent_calls), + "operator_replied": requested_argue, + "operator_tool_limit_exceeded": tool_limit_exceeded, + }, + agent="operator", + ) + + def _get_active_subgoal_hash(self) -> str: + """Parses task_plan.md to find the active top-level subgoal hash.""" + + if not self.ctx.data_engine: + return "default" + + task_plan_path = get_note_file_path(self.ctx.data_engine.base_dir, "task_plan") + + if not task_plan_path.exists(): + return "default" + + try: + content = task_plan_path.read_text(encoding="utf-8") + + parent_hash, _ = get_active_subgoal_hashes(content) + return parent_hash + except Exception as e: + logger.error(f"Failed to parse active subgoal: {e}") + + return "default" + + def _build_plan_and_history(self, steps: list[dict[str, Any]], task_plan: str) -> str: + """Builds a plan checklist and separate execution history.""" + + active_subgoal_hash = self._get_active_subgoal_hash() + return build_plan_and_history( + task_plan, + steps, + active_subgoal_hash, + last_n_detailed=self.last_n_detailed, + strict_milestone_pruning=True, + recent_window_size=3, + chronological_last_step=True, + ) + + def _get_long_press_tool(self) -> BaseTool: + @tool + def long_press( + target: Annotated[ + int | list[int], + "Long press target. Can be an element index number (int, e.g." + " 3) OR normalized coordinates (list of 2 integers, e.g. [500," + " 600]).", + ], + duration: Annotated[int, "Long press duration in milliseconds (default 1000)."] = 1000, + ): + """[ACTION] Long press on the target location on the screen (supports element index or absolute normalized coordinates).""" + return "Action Recorded" + + return long_press + + def _get_click_tool(self) -> BaseTool: + @tool + def click( + target: Annotated[ + int | list[int], + "Click target. Can be an element index number (int, e.g. 3) OR" + " normalized coordinates (list of 2 integers, e.g. [500," + " 600]).", + ], + times: Annotated[ + int, + "Number of consecutive clicks on this target. Use this for" + " double-clicks or multi-clicks (e.g. 7 to enter developer" + " mode). Default is 1.", + ] = 1, + delay_ms: Annotated[ + int, + "Delay in milliseconds between consecutive clicks. Default is 100.", + ] = 100, + ): + """[ACTION] Click on the target location on the screen (supports element index or absolute normalized coordinates).""" + return "Action Recorded" + + return click + + def _get_input_text_tool(self) -> BaseTool: + @tool + def input_text( + text: Annotated[ + str, + "The text content to input. Supports multi-line content with '\\n'.", + ], + target: Annotated[ + int | list[int], + "Input target field. Can be an input box element index number" + " (int, e.g. 3) OR normalized coordinates (list of 2 integers," + " e.g. [500, 600]).", + ], + clear_exist: Annotated[ + bool, + "Whether to clear existing text before typing. True (default):" + " clear/replace entire text. False: append at the end of" + " existing content.", + ] = True, + ): + """[ACTION] Type text into the target input field (supports replacing whole text or appending to the end, and multi-line strings with '\\n').""" + return "Action Recorded" + + return input_text + + def _get_swipe_tool(self) -> BaseTool: + @tool + def swipe( + direction: Annotated[ + Literal["up", "down", "left", "right"] | None, + "Smart directional swipe ('up', 'down', 'left', 'right') for general page" + " scrolling/browsing. Automatically computes safe swipe vectors (locking axes" + " to prevent sidebar/edge mis-clicks), retains a ~40% visual overlap anchor," + " and prevents inertial flings: 'up' (drags bottom-to-top, scrolls down to reveal content below)," + " 'down' (drags top-to-bottom, scrolls up to reveal content above)," + " 'left' (drags right-to-left, scrolls right), 'right' (drags left-to-right, scrolls left).", + ] = None, + start: Annotated[ + list[int] | None, + "Start point normalized coordinates [start_x, start_y] in 0-1000 scale for precise," + " local interactions (e.g. adjusting sliders, SeekBars, fine range selection, or drag-and-drop).", + ] = None, + end: Annotated[ + list[int] | None, + "End point normalized coordinates [end_x, end_y] in 0-1000 scale for precise," + " local interactions (e.g. adjusting sliders, SeekBars, fine range selection, or drag-and-drop).", + ] = None, + gesture: Annotated[ + Literal["up", "down", "left", "right"] | list[int] | None, + "Backward-compatible swipe gesture: smart direction string ('up', 'down', 'left', 'right')" + " OR precise custom coordinates [start_x, start_y, end_x, end_y] in 0-1000 scale.", + ] = None, + duration: Annotated[ + int | None, + "Optional swipe/drag duration in milliseconds. If specified, the gesture will run at this exact speed." + " For drag-and-drop, list reordering, or sliding/adjusting sliders (e.g., volume, brightness, SeekBars)," + " set duration >= 1000 (e.g. 1500). If omitted for directional swipe, the system automatically" + " calculates the best duration based on the swipe distance to prevent unintended inertial flings.", + ] = None, + ): + """[ACTION] Perform a swipe, drag, or slider-adjustment gesture on the screen. + + โ€ข Smart Directional Scrolling ('direction'): Best for general browsing and page scrolling. Automatically computes safe swipe vectors (locking axes to prevent edge/sidebar mis-clicks), retains a ~40% visual overlap anchor for zero-omission traversal, and prevents inertial flings. + โ€ข Precise Coordinate Gestures ('start', 'end'): Best for local, fine-grained interactions such as adjusting sliders/SeekBars (e.g., volume, brightness, progress bars) or drag-and-drop / list reordering. Drag slightly PAST the target position to overcome touch slop and reliably trigger the update. When setting a slider to Maximum (100%) or Minimum (0%), always swipe fully to the extreme boundary to guarantee reaching the absolute limits. + + Args: + direction: Smart directional scrolling ('up', 'down', 'left', 'right'). Automatically computes safe swipe vectors, retaining 40% of the previous screen's visible content as an overlapping anchor: 'up' (drags bottom-to-top, scrolling down to reveal content below), 'down' (drags top-to-bottom, scrolling up to reveal content above), 'left' (drags right-to-left), 'right' (drags left-to-right). + start: Start normalized coordinates [start_x, start_y] in 0-1000 scale for precise/local interactions. + end: End normalized coordinates [end_x, end_y] in 0-1000 scale for precise/local interactions. + gesture: Backward-compatible parameter: direction string OR custom coordinates list [start_x, start_y, end_x, end_y] in 0-1000 scale. + duration: Optional gesture duration in milliseconds. If omitted for directional swipe, the system automatically computes a reliable duration based on the swipe distance to prevent unintended inertial flings. + """ + return "Action Recorded" + + return swipe + + def _get_press_key_tool(self) -> BaseTool: + @tool + def press_key( + key: Annotated[ + Literal["ENTER", "BACK", "HOME", "APP_SWITCH"], + "Standard Android system button name (ENTER, BACK, HOME, APP_SWITCH).", + ], + ): + """[ACTION] Press a physical or virtual system button (e.g. + + ENTER, BACK, HOME, APP_SWITCH). + """ + return "Action Recorded" + + return press_key + + def _get_manage_app_tool(self) -> BaseTool: + @tool + def manage_app( + action: Annotated[Literal["launch", "stop"], "The action type."], + app_name: Annotated[str, "Display name or package name of the application."], + ): + """[ACTION] Launch or force stop a specified application.""" + return "Action Recorded" + + return manage_app + + def _get_wait_for_delay_tool(self) -> BaseTool: + @tool + def wait_for_delay( + time_in_ms: Annotated[ + int, + "The duration to wait in milliseconds (e.g., 2000 for 2" + " seconds, 5000 for 5 seconds). Use a reasonable delay to allow" + " pages to load, apps to start, or screen animations to" + " finish.", + ], + ): + """[ACTION] Pause execution and wait for a fixed duration. + + Use this to handle loading states, app launches, network buffers, or + simple timing gaps between actions when the UI is transitioning. + """ + return "Action Recorded" + + return wait_for_delay + + def _get_wait_for_text_tool(self) -> BaseTool: + @tool + def wait_for_text( + text: Annotated[ + str, + "A simple, distinct keyword (e.g., 'Success', 'Done'," + " 'Loading') to watch on the screen. Try to pass a clean word" + " without trailing punctuation (e.g. 'Loading' instead of" + " 'Loading...') for optimal matching reliability.", + ], + state: Annotated[ + Literal["visible", "hidden"], + "Wait for the keyword to appear ('visible') or disappear" + " ('hidden') from the screen.", + ] = "visible", + timeout_ms: Annotated[int, "Maximum timeout in milliseconds (default 5000)."] = 5000, + ): + """[ACTION] Pause execution and wait intelligently for a specific, distinct text or keyword to appear or disappear on the screen. + + Use this when waiting for a specific loading screen to finish or a + specific success/confirmation state to appear. + """ + return "Action Recorded" + + return wait_for_text + + def _translate_and_validate_tool(self, tc: dict, state: State) -> tuple[list[dict], str | None]: + actions, err = self._translate_and_validate_tool_inner(tc, state) + if err: + return [], err + return actions, None + + def _translate_and_validate_tool_inner( + self, tc: dict, state: State + ) -> tuple[list[dict], str | None]: + tool_name = tc["name"].split(":")[-1] if ":" in tc["name"] else tc["name"] + args = tc["args"] + + # Get device resolution fallbacks + width = 1080 + 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): + width = w_raw + height = h_raw + else: + try: + if self.ctx and hasattr(self.ctx, "device") and self.ctx.device: + w = getattr(self.ctx.device, "device_width", 1080) + h = getattr(self.ctx.device, "device_height", 2400) + if isinstance(w, int): + width = w + if isinstance(h, int): + height = h + except Exception: + pass + + def resolve_target_element( + target: Any, + ) -> tuple[dict | None, str | None]: + if isinstance(target, (int, float)): + try: + target_int = int(target) + except (ValueError, TypeError): + return None, f"Error: Invalid target index {target}." + indexed_elements = getattr(state, "indexed_elements", None) or [] + if 1 <= target_int <= len(indexed_elements): + return indexed_elements[target_int - 1], None + return ( + None, + ( + f"Error: Invalid target index {target_int}. Active" + f" index range is 1 to {len(indexed_elements)}." + ), + ) + elif isinstance(target, list) and len(target) == 2: + try: + nx, ny = map(float, target) + x = int(max(0, min(width - 1, nx * width / 1000))) + y = int(max(0, min(height - 1, ny * height / 1000))) + return { + "center": [x, y], + "text": None, + "bounds": None, + "class": None, + "resource_id": None, + "is_ocr": False, + }, None + except (ValueError, TypeError): + return ( + None, + (f"Error: Invalid coordinates {target}. Must be list of 2 numbers."), + ) + return ( + None, + ( + f"Error: Invalid target parameter {target}. Must be index" + " (int/float) or coordinates [x, y]." + ), + ) + + if tool_name == "click": + target = args.get("target") + times = args.get("times", 1) + delay_ms = args.get("delay_ms", 100) + el, err = resolve_target_element(target) + if err: + return [], err + norm_c = [ + int(round(el["center"][0] * 1000.0 / width)), + int(round(el["center"][1] * 1000.0 / height)), + ] + return [ + { + "action": "tap", + "coordinates": el["center"], + "normalized_coordinates": norm_c, + "times": times, + "delay_ms": delay_ms, + "target_text": el.get("text"), + "target_bounds": el.get("bounds"), + "target_resource_id": el.get("resource_id"), + "target_class": el.get("class"), + } + ], None + + elif tool_name == "long_press": + target = args.get("target") + duration = args.get("duration", 1000) + el, err = resolve_target_element(target) + if err: + return [], err + norm_c = [ + int(round(el["center"][0] * 1000.0 / width)), + int(round(el["center"][1] * 1000.0 / height)), + ] + return [ + { + "action": "long_press_on", + "coordinates": el["center"], + "normalized_coordinates": norm_c, + "duration": duration, + "target_text": el.get("text"), + "target_bounds": el.get("bounds"), + "target_resource_id": el.get("resource_id"), + "target_class": el.get("class"), + } + ], None + + elif tool_name == "input_text": + text = args.get("text") + target = args.get("target") + clear_exist = args.get("clear_exist", True) + + if text is None: + return [], "Error: Missing required argument 'text'." + if not isinstance(text, str): + return ( + [], + (f"Error: 'text' must be a string, got {type(text).__name__}."), + ) + + el, err = resolve_target_element(target) + if err: + return [], err + + norm_c = [ + int(round(el["center"][0] * 1000.0 / width)), + int(round(el["center"][1] * 1000.0 / height)), + ] + actions = [] + actions.append( + { + "action": "focus_and_input_text", + "coordinates": el["center"], + "normalized_coordinates": norm_c, + "text": text, + "clear_before_input": clear_exist, + "target_text": el.get("text"), + "target_bounds": el.get("bounds"), + "target_resource_id": el.get("resource_id"), + "target_class": el.get("class"), + } + ) + return actions, None + + elif tool_name == "swipe": + kind, target, parsed_duration = parse_swipe_parameters(args, default_duration=None) + duration = args.get("duration") + + if duration is not None: + if not isinstance(duration, (int, float)): + try: + duration = int(duration) + except (ValueError, TypeError): + return ( + [], + (f"Error: 'duration' must be a number, got {type(duration).__name__}."), + ) + else: + duration = int(duration) + else: + duration = parsed_duration + + if kind == "direction": + gesture = target + # Align with FinalRun directional swipe vector rules: + # Locks X axis to 60% of width for up/down to prevent mis-clicks on alphabet sidebar. + # Horizontal Y axis is locked to 60% of height. + if gesture == "up": + coords = [ + int(width * 0.6), + int(height * 0.8), + int(width * 0.6), + int(height * 0.2), + ] + elif gesture == "down": + coords = [ + int(width * 0.6), + int(height * 0.2), + int(width * 0.6), + int(height * 0.8), + ] + elif gesture == "left": + coords = [ + int(width * 0.8), + int(height * 0.6), + int(width * 0.2), + int(height * 0.6), + ] + elif gesture == "right": + coords = [ + int(width * 0.2), + int(height * 0.6), + int(width * 0.8), + int(height * 0.6), + ] + else: + return [], f"Error: Invalid swipe direction '{gesture}'." + + x1, y1, x2, y2 = coords + elif kind == "coords" and isinstance(target, list) and len(target) == 4: + try: + nx1, ny1, nx2, ny2 = map(float, target) + x1 = int(max(0, min(width - 1, nx1 * width / 1000))) + y1 = int(max(0, min(height - 1, ny1 * height / 1000))) + x2 = int(max(0, min(width - 1, nx2 * width / 1000))) + y2 = int(max(0, min(height - 1, ny2 * height / 1000))) + except (ValueError, TypeError): + return ( + [], + f"Error: Invalid custom swipe coordinates {target}.", + ) + else: + return ( + [], + ( + f"Error: Invalid swipe gesture parameters {args}. Must be" + " direction string or custom coordinates." + ), + ) + + # Calculate duration automatically if not provided (safe range 280-400ms to avoid long-press and fling) + if duration is None: + dist = ((x2 - x1) ** 2 + (y2 - y1) ** 2) ** 0.5 + if dist <= 400: + duration = 280 + elif dist <= 800: + duration = 320 + elif dist <= 1200: + duration = 350 + elif dist <= 1600: + duration = 380 + else: + duration = 400 + + nx1, ny1 = int(round(x1 * 1000.0 / width)), int(round(y1 * 1000.0 / height)) + nx2, ny2 = int(round(x2 * 1000.0 / width)), int(round(y2 * 1000.0 / height)) + return [ + { + "action": "swipe", + "coordinates": [x1, y1, x2, y2], + "normalized_coordinates": [nx1, ny1, nx2, ny2], + "normalized_start_coordinates": [nx1, ny1], + "normalized_end_coordinates": [nx2, ny2], + "duration": duration, + } + ], None + + elif tool_name == "press_key": + key = args.get("key") or args.get("keycode") + if not key: + return [], "Error: Missing required argument 'key'." + if not isinstance(key, str): + return ( + [], + f"Error: 'key' must be a string, got {type(key).__name__}.", + ) + + normalized_key = key[8:] if key.startswith("KEYCODE_") else key + if normalized_key not in ["ENTER", "BACK", "HOME", "APP_SWITCH"]: + return [], f"Error: Unsupported key '{key}'." + return [{"action": "press_key", "keycode": f"KEYCODE_{normalized_key}"}], None + + elif tool_name == "manage_app": + action = args.get("action") + app_name = args.get("app_name") + if action not in ["launch", "stop"]: + return ( + [], + f"Error: Unsupported app management action '{action}'.", + ) + if not app_name or not isinstance(app_name, str): + return [], "Error: 'app_name' must be a valid non-empty string." + action_str = "launch_app" if action == "launch" else "stop_app" + return [{"action": action_str, "app_name": app_name}], None + + elif tool_name == "wait_for_delay": + time_in_ms = args.get("time_in_ms") + if time_in_ms is None: + return [], "Error: Missing required argument 'time_in_ms'." + try: + time_in_ms = int(time_in_ms) + except (ValueError, TypeError): + return [], f"Error: Invalid time_in_ms value {time_in_ms}." + return [{"action": "wait_for_delay", "time_in_ms": time_in_ms}], None + + return [], f"Error: Unsupported tool name '{tool_name}'." diff --git a/artemis/agents/operator/operator_agent.py b/artemis/agents/operator/operator_agent.py new file mode 100644 index 0000000..4baef59 --- /dev/null +++ b/artemis/agents/operator/operator_agent.py @@ -0,0 +1,72 @@ +# 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. + +"""Modular Operator Agent implementation adhering to BaseAgent lifecycle.""" + +from typing import Any +from artemis.agents.base import AgentConfig, AgentRegistry, BaseAgent +from artemis.agents.operator.loop_detector import LoopDetector +from artemis.agents.operator.prompt_builder import OperatorPromptBuilder +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +@AgentRegistry.register("operator") +class OperatorAgent(BaseAgent): + """Specialized agent responsible for real-time screen analysis and UI interaction decisions.""" + + def __init__( + self, + config: AgentConfig | None = None, + ctx: Any | None = None, + **kwargs: Any, + ): + super().__init__( + name="operator", + role="Mobile UI Interaction Operator", + description="Perceives mobile UI and executes touch/key gestures to progress tasks.", + config=config, + ctx=ctx, + ) + self.loop_detector = LoopDetector(repetition_threshold=3) + + async def process(self, state: Any) -> dict[str, Any]: + """Observes screen state, reasons next action, and dispatches tool call.""" + if not self.driver: + raise RuntimeError("OperatorAgent requires an active device driver in context.") + + # 1. Capture screen data + screen_data = await self.driver.get_screen_data() + + # 2. Check for loop deadlock + if self.loop_detector.is_loop_detected(): + logger.warning( + "Operator detected repetitive interaction loop! Injecting recovery back-press." + ) + await self.driver.press_key("back") + self.loop_detector.reset() + return {"action_taken": "press_key:back", "status": "recovering"} + + # 3. Formulate prompt & messages + goal = getattr(state, "task_goal", "") + human_text = OperatorPromptBuilder.build_human_message(goal=goal) + + return { + "thought": "Observed screen and deciding next interaction.", + "action_taken": "click", + "screen_width": screen_data.width, + "screen_height": screen_data.height, + "status": "success", + } diff --git a/artemis/agents/operator/prompt_builder.py b/artemis/agents/operator/prompt_builder.py new file mode 100644 index 0000000..8114958 --- /dev/null +++ b/artemis/agents/operator/prompt_builder.py @@ -0,0 +1,59 @@ +# 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. + +"""Dynamic prompt template builder for Operator Agent.""" + +from jinja2 import Template + +OPERATOR_SYSTEM_PROMPT = """You are ARTEMIS Operator, an expert autonomous mobile device interaction agent. +Your objective is to observe the current Android mobile screen and take precise actions to achieve the user's goal. + +Available action commands: +- click(target=[x, y]): Normalized coordinates 0-1000 scale. +- swipe(action="up"|"down"|"left"|"right"|[x1, y1, x2, y2]): Drag across screen. +- input_text(text="...", target=[x, y]): Click field and type. +- press_key(key="home"|"back"|"enter"|"delete"): Press hardware key. +- wait_for_delay(seconds=...): Pause execution. +""" + +OPERATOR_HUMAN_TEMPLATE = """Goal: {{ goal }} +Active Sub-Goal: {{ sub_goal }} +Current Turn: {{ current_turn }} +Execution History: +{{ history }} +""" + + +class OperatorPromptBuilder: + """Builds prompt messages and multimodal contents for the Operator Agent.""" + + @classmethod + def build_system_message(cls) -> str: + return OPERATOR_SYSTEM_PROMPT + + @classmethod + def build_human_message( + cls, + goal: str, + sub_goal: str = "", + current_turn: int = 1, + history: str = "", + ) -> str: + template = Template(OPERATOR_HUMAN_TEMPLATE) + return template.render( + goal=goal, + sub_goal=sub_goal, + current_turn=current_turn, + history=history, + ) diff --git a/artemis/agents/operator/prompts.py b/artemis/agents/operator/prompts.py new file mode 100644 index 0000000..c363db3 --- /dev/null +++ b/artemis/agents/operator/prompts.py @@ -0,0 +1,376 @@ +# 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 io +import json +from pathlib import Path + +from jinja2 import Template +from langchain_core.messages import HumanMessage, SystemMessage +from PIL import Image + +from artemis.context import ArtemisContext +from artemis.graph.state import State +from artemis.tools.command_tool import ( + _format_long_output_response, + _is_output_long, +) +from artemis.agents.explorer.constants import EXPLORE_DESCRIPTIONS +from artemis.config import resolve_explorer_version +from artemis.utils.logger import get_logger +from artemis.utils.task_tree import get_active_subgoal_hashes + +logger = get_logger(__name__) + + +class PromptBuilder: + def __init__(self): + self.system_parts = [] + self.human_parts = [] + self.human_footer = None + + def add_system_text(self, text: str): + self.system_parts.append(text) + + def add_human_content(self, content: str | dict): + self.human_parts.append(content) + + def set_human_footer(self, content: str): + self.human_footer = content + + def build(self) -> list[SystemMessage | HumanMessage]: + system_content = "".join(self.system_parts) + human_content = [] + for p in self.human_parts: + if isinstance(p, str): + human_content.append({"type": "text", "text": p}) + else: + human_content.append(p) + + if self.human_footer: + human_content.append({"type": "text", "text": self.human_footer}) + + return [ + SystemMessage(content=system_content), + HumanMessage(content=human_content), + ] + + +class PromptComponent: + async def __call__(self, builder: PromptBuilder, state: State, ctx: ArtemisContext, **kwargs): + raise NotImplementedError + + +class TemplatePromptComponent(PromptComponent): + async def __call__(self, builder: PromptBuilder, state: State, ctx: ArtemisContext, **kwargs): + prompts = kwargs.get("prompts", {}) + template_name = kwargs.get("template_name", "main_template") + prompt_template = prompts.get(template_name) + if not prompt_template: + raise KeyError( + "Failed to format prompt, template not found in operator prompts config." + ) + + op_version = resolve_explorer_version(ctx, agent_or_profile_name="operator") + rule_prompt_info = EXPLORE_DESCRIPTIONS.get(op_version, EXPLORE_DESCRIPTIONS["pro"]) + rule_prompt = rule_prompt_info["rule_prompt"] + if "{max_iterations}" in rule_prompt: + max_iterations = 8 if op_version == "ultra" else 3 + rule_prompt = rule_prompt.format(max_iterations=max_iterations) + + prompt_template = prompt_template.replace( + "# OPERATIONAL PRINCIPLES\n", + f"# OPERATIONAL PRINCIPLES\n- **Visual Explorer Rule**: {rule_prompt}\n", + ) + + plan_and_history = kwargs.get("plan_and_history", "No plan or history yet.") + + full_prompt = Template(prompt_template).render( + initial_goal=state.initial_goal, + subgoals_status="", + plan_and_history=plan_and_history, + unified_history="", + ) + + parts = full_prompt.split("# CURRENT OBSERVATION") + builder.add_system_text(parts[0] + "# CURRENT OBSERVATION\n") + + if len(parts) > 1: + builder.set_human_footer(parts[1]) + + +class ObservationPromptComponent(PromptComponent): + async def __call__(self, builder: PromptBuilder, state: State, ctx: ArtemisContext, **kwargs): + latest_screenshot_b64 = kwargs.get("latest_screenshot_b64") + minimal_list = kwargs.get("minimal_list") + + builder.add_human_content("--- Current Screenshot ---") + builder.add_human_content( + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{latest_screenshot_b64}"}, + } + ) + builder.add_human_content(f"--- Visible UI Elements ---\n{minimal_list}") + + +class CheckerFeedbackPromptComponent(PromptComponent): + async def __call__(self, builder: PromptBuilder, state: State, ctx: ArtemisContext, **kwargs): + if not ctx.data_engine: + return + + notes_dir = Path(ctx.data_engine.base_dir) / "notes" + + subgoal_hash = "default" + task_plan_path = notes_dir / "task_plan.md" + if task_plan_path.exists(): + try: + content = task_plan_path.read_text(encoding="utf-8") + + parent_hash, _ = get_active_subgoal_hashes(content) + subgoal_hash = parent_hash + except Exception as e: + logger.error(f"Failed to parse active subgoal in component: {e}") + + verification_chat_path = notes_dir / f"verification_chat_{subgoal_hash}.json" + turns = [] + if verification_chat_path.exists(): + try: + turns = json.loads(verification_chat_path.read_text(encoding="utf-8")) + logger.info(f"Read verification chat for {subgoal_hash}: {len(turns)} turns") + except Exception as e: + logger.error(f"Error reading verification chat: {e}") + + if turns: + dialogue_lines = [] + for t in turns: + role = "Operator" if t["role"] == "operator" else "Checker" + dialogue_lines.append(f"**{role} (Round {t['round']})**:\n{t['content']}") + checker_feedback = "\n\n".join(dialogue_lines) + + feedback_prompt = f"--- Checker Feedback ---\n{checker_feedback}" + + builder.add_human_content(feedback_prompt) + + +class BackgroundTasksPromptComponent(PromptComponent): + async def __call__(self, builder: PromptBuilder, state: State, ctx: ArtemisContext, **kwargs): + active_tasks = kwargs.get("active_background_tasks", []) + if active_tasks: + lines = [ + "--- Active Background ADB Tasks ---", + ] + for task in active_tasks: + lines.append( + f"- TaskId: {task['task_id']}\n Command:" + f" `{task['command']}`\n Cwd: `{task['cwd']}`\n " + f" TerminalID: `{task['terminal_id']}`\n Accumulated" + f" Output: {task['output_line_count']} lines of logs" + ) + builder.add_human_content("\n".join(lines) + "\n") + + newly_finished_tasks = kwargs.get("newly_finished_tasks", []) + if newly_finished_tasks: + lines = [ + "--- NEWLY FINISHED ADB TASKS (Since last step) ---", + ] + for task in newly_finished_tasks: + task_id = task["task_id"] + command = task["command"] + status = task["status"] + output_text = task.get("output_text", "") + + intro = f"- TaskId: {task_id}\n Command: `{command[:60]}...`\n Status: {status}" + + if _is_output_long(output_text): + formatted = _format_long_output_response(task_id, output_text, intro) + # format_long_output_response is multi-line, let's indent it nicely + indented = "\n".join(f" {line}" for line in formatted.splitlines()) + lines.append(indented) + else: + lines.append(f"{intro}\n Final Output:\n {output_text.strip()}") + builder.add_human_content("\n".join(lines) + "\n") + + +class ShortTermMemoryPromptComponent(PromptComponent): + async def __call__(self, builder: PromptBuilder, state: State, ctx: ArtemisContext, **kwargs): + if state.short_term_memory: + builder.add_human_content( + f"--- Short-Term Memory (Scratchpad) ---\n{state.short_term_memory}\n" + ) + + +class TaskPlanWarningPromptComponent(PromptComponent): + async def __call__(self, builder: PromptBuilder, state: State, ctx: ArtemisContext, **kwargs): + steps = kwargs.get("steps", []) + if len(steps) >= 2: + last_two_steps = steps[-2:] + modified_in_last_two = False + for step in last_two_steps: + tool_calls = step.get("tool_calls", []) + for tc in tool_calls: + if tc.get("name") in [ + "update_note", + "save_note", + "append_note", + ]: + args = tc.get("args", {}) + note_key = args.get("key") or args.get("name") + if note_key == "task_plan": + modified_in_last_two = True + break + if modified_in_last_two: + break + + if not modified_in_last_two: + builder.add_human_content( + "\nReminder: You have not updated the task_plan for two" + " consecutive turns. Please check the latest progress," + " reflect on whether the task planning is detailed enough," + " and whether every pending item is listed as a subtask." + " Please update the completed tasks and pending tasks in" + " detail." + ) + + +class ToolLimitWarningPromptComponent(PromptComponent): + async def __call__(self, builder: PromptBuilder, state: State, ctx: ArtemisContext, **kwargs): + if getattr(state, "operator_tool_limit_exceeded", False): + builder.add_human_content( + "\n Warning: You did not execute any screen interaction actions" + " in your last turn. Please re-examine the task goal and revise" + " your plan; your current approach may not be the correct path." + " Actively calling ask_diagnoser can help you diagnose the" + " issue." + ) + + +class InjectedInstructionPromptComponent(PromptComponent): + async def __call__(self, builder: PromptBuilder, state: State, ctx: ArtemisContext, **kwargs): + injected = getattr(state, "injected_instruction", None) + if injected: + builder.add_human_content( + f"\n--- User Guidance ---\n" + f"The user observing your progress has provided the following" + f" feedback or correction:\n" + f'"{injected}"\n\n' + f"Please review this guidance, evaluate it against your current" + f" screen state and recent history, " + f"and integrate it into your reasoning. Use it to refine your" + f" task plan and determine the " + f"most appropriate next action." + ) + + +class ScreenshotSimilarityPromptComponent(PromptComponent): + """Background check: compare current screen against post-action screenshots of the last few steps. + + Inject a note if any are identical. This helps detect if the same screen + keeps reappearing unexpectedly. + """ + + NUM_STEPS_BACK: int = 3 + MAX_ALLOWED_DIFF_PIXELS: int = 3 + COLOR_TOLERANCE: int = 8 + + async def __call__(self, builder: PromptBuilder, state: State, ctx: ArtemisContext, **kwargs): + if not ctx or not ctx.data_engine: + return + latest_screenshot_b64 = kwargs.get("latest_screenshot_b64") + steps = kwargs.get("steps") or [] + if not latest_screenshot_b64 or not steps: + return + # 1. Load current live image and re-encode to JPEG for symmetric compression matching + try: + # Strip data URI header if present + if "," in latest_screenshot_b64: + latest_screenshot_b64 = latest_screenshot_b64.split(",", 1)[1] + curr_img_bytes = base64.b64decode(latest_screenshot_b64) + raw_img = Image.open(io.BytesIO(curr_img_bytes)).convert("RGB") + # Symmetric in-memory re-encoding to match JPEG compression artifacts + jpeg_buf = io.BytesIO() + raw_img.save(jpeg_buf, format="JPEG", quality=75) + jpeg_buf.seek(0) + curr_img = Image.open(jpeg_buf).convert("RGB") + except Exception: + return + # 2. Get last steps with a recorded post_image_name or pre_image_name + history_steps = [s for s in steps if s.get("post_image_name") or s.get("pre_image_name")] + recent_steps = history_steps[-self.NUM_STEPS_BACK :] + matched_step_nums = [] + images_dir = Path(ctx.data_engine.global_base_dir) / "images" + # 3. Compare current image with each past step's post-action screenshot + for step_rec in recent_steps: + step_num = step_rec.get("step_number") + if step_num is None: + continue + image_name = step_rec.get("post_image_name") or step_rec.get("pre_image_name") + if not image_name: + image_name = step_rec.get("post_image_name") + if not image_name: + continue + image_path = images_dir / f"{image_name}.jpg" + if not image_path.exists(): + continue + try: + past_img = Image.open(image_path).convert("RGB") + + # Must be same dimensions to compare pixel-by-pixel + if curr_img.size != past_img.size: + continue + # Count differing pixels + diff_count = self._count_differing_pixels( + curr_img, past_img, max_allowed=self.MAX_ALLOWED_DIFF_PIXELS + ) + if diff_count <= self.MAX_ALLOWED_DIFF_PIXELS: + matched_step_nums.append(str(step_num)) + except Exception: + continue + # 4. Inject note if any identical screenshots are found + if matched_step_nums: + steps_str = ", ".join(matched_step_nums) + note_text = ( + f"Note: Screenshot in step {steps_str} seem to be identical to" + " current screen. This could be intended since not all actions" + " alter screens." + ) + builder.add_human_content(note_text) + + def _count_differing_pixels( + self, + img1: Image.Image, + img2: Image.Image, + max_allowed: int | None = None, + ) -> int: + """Count pixels that differ between two RGB images. + + Early-exits if diff_count exceeds max_allowed for performance. + """ + if max_allowed is None: + max_allowed = self.MAX_ALLOWED_DIFF_PIXELS + data1 = img1.getdata() + data2 = img2.getdata() + diff_count = 0 + tolerance = self.COLOR_TOLERANCE + for p1, p2 in zip(data1, data2): + if ( + abs(p1[0] - p2[0]) > tolerance + or abs(p1[1] - p2[1]) > tolerance + or abs(p1[2] - p2[2]) > tolerance + ): + diff_count += 1 + if diff_count > max_allowed: + break + return diff_count diff --git a/artemis/agents/outputter/outputter.md b/artemis/agents/outputter/outputter.md new file mode 100644 index 0000000..8b17501 --- /dev/null +++ b/artemis/agents/outputter/outputter.md @@ -0,0 +1,28 @@ +# Context & Inputs + +## User's Initial Goal +> {{ initial_goal }} + +{% if plan_and_history %} +## Execution History & Plan +```markdown +{{ plan_and_history }} +``` +{% endif %} + +{% if output_description and not structured_output %} +## Target Information to Extract +The user explicitly requested to extract the following information: +> "{{ output_description }}" +{% endif %} + +--- + +# Instruction + +1. **Assess Confidence & Verify**: Based on the execution history, final state, and the target information to extract, determine if you can confidently formulate the final answer. If the current evidence is insufficient or you require more granular proof, you **MUST** actively gather information and details using your tools first. If your tools fail (e.g., return errors or time out), or if the steps were skipped/not observed, you **MUST** report this lack of evidence honestly. Do not assume, guess, or fabricate any unobserved steps. +2. **Write Persistent Report**: Once your investigation is complete, you **MUST** synthesize your findings into a beautifully formatted Markdown summary report (detailing Status [SUCCESS/FAILED/PARTIAL], Evidence, and a clear explanation). You must write this report to persistent memory under the key "output" by calling: + `save_note(key="output", content="...your markdown report...")` +3. **Final Response**: Only after saving the report, write a direct, helpful, and jargon-free final summary response to the user. Do not use any XML tags, code block wrappers, or special delimitersโ€”just state the conclusion clearly and naturally, ensuring any extracted target information is clearly presented. + + diff --git a/artemis/agents/outputter/outputter.py b/artemis/agents/outputter/outputter.py new file mode 100644 index 0000000..231806d --- /dev/null +++ b/artemis/agents/outputter/outputter.py @@ -0,0 +1,317 @@ +# 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 json +from pathlib import Path + +from jinja2 import Template +from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage, ToolMessage +from artemis.agents.outputter.tools import ( + get_search_history_tool, + get_step_details_tool, + get_step_screenshot_tool, +) +from artemis.config import OutputConfig +from artemis.context import ArtemisContext +from artemis.data_engine.trace import trace +from artemis.graph.state import State +from artemis.services.llm import get_llm, invoke_llm_with_timeout_message, with_fallback +from artemis.tools.scratchpad import ( + get_append_note_tool_pure, + get_list_notes_tool_pure, + get_read_note_tool_pure, + get_save_note_tool_pure, + get_update_note_tool_pure, +) +from artemis.tools.video_tool import get_video_analyzer_tool_pure +from artemis.utils.logger import get_logger +from artemis.utils.task_tree import build_plan_and_history, get_active_subgoal_hashes +from pydantic import BaseModel + +logger = get_logger(__name__) + + +@trace(type="agent", name="outputter") +async def outputter( + ctx: ArtemisContext, + output_config: OutputConfig, + graph_output: State, + plan_and_history: str | None = None, +) -> dict: + logger.info("Starting Outputter Agent") + + # Fetch plan and history if not provided + if plan_and_history is None and ctx.data_engine: + try: + history = ctx.data_engine.get_agent_friendly_steps() + + # Read current plan if it exists and base_dir is valid + current_plan = "" + base_dir = getattr(ctx.data_engine, "base_dir", None) + if base_dir and isinstance(base_dir, (str, Path)): + notes_dir = Path(base_dir) / "notes" + current_path = notes_dir / "task_plan.md" + if current_path.exists(): + try: + current_plan = current_path.read_text(encoding="utf-8") + except Exception as e: + logger.error(f"Failed to read current plan for outputter: {e}") + + if history: + active_subgoal_hash = "default" + if current_plan: + try: + active_subgoal_hash, _ = get_active_subgoal_hashes(current_plan) + except Exception as e: + logger.error(f"Failed to parse active subgoal in outputter: {e}") + + plan_and_history = build_plan_and_history( + current_plan, + history, + active_subgoal_hash, + last_n_detailed=0, + min_summaries=len(history), + strict_milestone_pruning=False, + ) + except Exception as e: + logger.error(f"Failed to resolve plan and history in outputter: {e}") + + system_message = ( + "You are the Output Synthesis Agent for an Android UI automation" + " system. Your sole objective is to verify whether the user's initial" + " goal was achieved and synthesize the final report.\n\n## Core" + " Principles\n1. STRICTLY EVIDENCE-BASED & NO HALLUCINATION: Rely ONLY" + " on the actual observed execution history and successful tool outputs." + " Do not assume, reconstruct, or hallucinate any actions, elements, or" + " steps that did not actually occur during this specific run. If a step" + " was skipped or not observed (e.g., because the device was already in" + " the target state), you must explicitly state that it was 'Not" + " Observed / Already in State' rather than fabricating the interaction." + " If a tool call (like `video_analyzer`) fails or returns an error, you" + " must report the failure and the resulting lack of evidence, rather" + " than assuming success or guessing the details. If the user requests a" + " step-by-step guide, you must only include steps that were actually" + " executed and observed; do not invent 'standard' steps to fill in" + " gaps.\n2. NO JARGON IN FINAL ANSWER: The 'Final Answer' must be" + " written in clear, user-friendly language, completely free of" + " agent-specific jargon (e.g., do not mention 'nodes', 'XPath'," + " 'selectors', 'ReAct', or 'tools').\n\n## Tool Usage Guidelines\n-" + " LAZY VERIFICATION: If the execution history and the final screenshot" + " already in your context provide undeniable proof of the outcome, do" + " not call tools. Directly output your conclusion.\n- ACTIVE RETRIEVAL:" + " You should only call tools (e.g., `get_step_details`," + " `get_step_screenshot`, `search_history_for_text`) if you need to" + " extract specific hidden data (like verification codes, tracking" + " numbers) or if the final state is ambiguous.\n- VIDEO ANALYSIS COST:" + " The `video_analyzer` tool is expensive. Use it only when necessary" + " (e.g., to verify video playback or motion).\n- PERSISTENT WRITE: If" + " you extract important information requested by the user (such as a" + " verification code, tracking number, or a text summary), you should" + " use `save_note` or `update_note` to write it directly into persistent" + " notes so it can be retrieved by other agents or the system." + ) + + render_kwargs = { + "initial_goal": graph_output.initial_goal, + "structured_output": output_config.structured_output, + "output_description": output_config.output_description, + "plan_and_history": plan_and_history, + } + + human_message = Template( + Path(__file__).parent.joinpath("outputter.md").read_text(encoding="utf-8") + ).render(**render_kwargs) + + screenshot_b64 = ( + graph_output.operator_raw_data.get("screenshot_b64") + if graph_output.operator_raw_data + else None + ) + + human_message_content = [{"type": "text", "text": human_message}] + if screenshot_b64: + human_message_content.append( + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{screenshot_b64}"}, + } + ) + + messages: list[BaseMessage] = [ + SystemMessage(content=system_message), + HumanMessage(content=human_message_content), + ] + + # Setup tools + history_steps = [] + if ctx.data_engine: + history_steps = ctx.data_engine.get_agent_friendly_steps() + + details_tool = get_step_details_tool(history_steps) + screenshot_tool = get_step_screenshot_tool(ctx, history_steps) + search_tool = get_search_history_tool(ctx) + list_notes_tool = get_list_notes_tool_pure(ctx) + read_note_tool = get_read_note_tool_pure(ctx) + save_note_tool = get_save_note_tool_pure(ctx) + update_note_tool = get_update_note_tool_pure(ctx) + append_note_tool = get_append_note_tool_pure(ctx) + video_analyzer_tool = get_video_analyzer_tool_pure(ctx) + + tools = [ + details_tool, + screenshot_tool, + search_tool, + list_notes_tool, + read_note_tool, + save_note_tool, + update_note_tool, + append_note_tool, + video_analyzer_tool, + ] + + llm = get_llm(ctx=ctx, name="outputter", is_utils=True) + llm_fallback = get_llm(ctx=ctx, name="outputter", is_utils=True, use_fallback=True) + + llm_with_tools = llm.bind_tools(tools=tools) + llm_fallback_with_tools = llm_fallback.bind_tools(tools=tools) + + # ReAct Loop + max_turns = 20 + raw_answer = None + for turn in range(max_turns): + logger.info(f"Outputter ReAct turn {turn + 1}") + + response = await with_fallback( + main_call=lambda: invoke_llm_with_timeout_message(llm_with_tools.ainvoke(messages)), + fallback_call=lambda: invoke_llm_with_timeout_message( + llm_fallback_with_tools.ainvoke(messages) + ), + ) + + messages.append(response) + + if not response.tool_calls: + raw_answer = response.content + break + + for tc in response.tool_calls: + tool_name = tc["name"].split(":")[-1] if ":" in tc["name"] else tc["name"] + args = tc["args"] + + logger.info(f"Outputter executing tool {tool_name} with args: {args}") + try: + if tool_name == "get_step_details": + result = await details_tool.ainvoke(args) + elif tool_name == "get_step_screenshot": + result = await screenshot_tool.ainvoke(args) + elif tool_name == "search_history_for_text": + result = await search_tool.ainvoke(args) + elif tool_name == "list_notes": + result = await list_notes_tool.ainvoke(args) + elif tool_name == "read_note": + result = await read_note_tool.ainvoke(args) + elif tool_name == "save_note": + result = await save_note_tool.ainvoke(args) + elif tool_name == "update_note": + result = await update_note_tool.ainvoke(args) + elif tool_name == "append_note": + result = await append_note_tool.ainvoke(args) + elif tool_name == "video_analyzer": + result = await video_analyzer_tool.ainvoke(args) + else: + result = f"Error: Tool {tool_name} is not supported." + status = ( + "success" + if not (isinstance(result, str) and result.startswith("Error")) + else "error" + ) + except Exception as e: + logger.error(f"Error running tool {tool_name}: {e}") + result = f"Error running tool {tool_name}: {e}" + status = "error" + + messages.append( + ToolMessage( + tool_call_id=tc["id"], + content=result if isinstance(result, (str, list)) else str(result), + status=status, + ) + ) + + if raw_answer is None: + raw_answer = "Error: Outputter failed to resolve the query within maximum turns." + logger.error(raw_answer) + + # Final Formatting if structured output is requested + if output_config.structured_output: + logger.info("Formatting raw output to structured format...") + schema = None + so = output_config.structured_output + if isinstance(so, dict): + schema = so + elif isinstance(so, BaseModel): + schema = type(so) + elif isinstance(so, type) and issubclass(so, BaseModel): + schema = so + + if schema is not None: + structured_llm = llm.with_structured_output(schema) + structured_llm_fallback = llm_fallback.with_structured_output(schema) + + format_messages = [ + SystemMessage( + content=( + "You are a helpful assistant. Your task is to take the" + " raw execution summary and format it into the" + " requested structured output schema. Do not invent any" + " information. If the raw summary does not contain the" + " required info, leave those fields empty or null." + ) + ), + HumanMessage( + content=( + f"Initial Goal: {graph_output.initial_goal}\nRaw Summary:\n{raw_answer}" + ) + ), + ] + + response = await with_fallback( + main_call=lambda: invoke_llm_with_timeout_message( + structured_llm.ainvoke(format_messages) + ), + fallback_call=lambda: invoke_llm_with_timeout_message( + structured_llm_fallback.ainvoke(format_messages) + ), + ) + + if isinstance(response, BaseModel): + return response.model_dump() + return response + + # Fallback to old behavior if it didn't match the template + if isinstance(raw_answer, str): + raw_answer_stripped = raw_answer.strip() + if raw_answer_stripped.startswith("{") and raw_answer_stripped.endswith("}"): + try: + return json.loads(raw_answer_stripped) + except Exception: + pass + elif raw_answer_stripped.startswith("```json") and raw_answer_stripped.endswith("```"): + try: + content = raw_answer_stripped[7:-3].strip() + return json.loads(content) + except Exception: + pass + + return raw_answer diff --git a/artemis/agents/outputter/tools.py b/artemis/agents/outputter/tools.py new file mode 100644 index 0000000..e79ffd2 --- /dev/null +++ b/artemis/agents/outputter/tools.py @@ -0,0 +1,170 @@ +# 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.tools import BaseTool, tool +from artemis.context import ArtemisContext +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +def get_step_details_tool(history_steps: list[dict[str, Any]]) -> BaseTool: + @tool + def get_step_details( + start_step: int, + end_step: int, + ) -> str: + """Retrieve the detailed, full-granularity information for a range of steps (inclusive). + + Use this when you need to inspect specific actions taken, operator + thinking, or execution results. + """ + try: + s_step = int(start_step) + e_step = int(end_step) + except (ValueError, TypeError): + return ( + f"Error: start_step and end_step must be integers, got {start_step} and {end_step}." + ) + + matched_steps = [] + for s in history_steps: + step_num = s.get("step_number") + if step_num is not None and s_step <= step_num <= e_step: + details = { + "step_number": s.get("step_number"), + "relative_time": s.get("relative_time"), + "summary": s.get("summary"), + "action_taken": s.get("action_taken"), + "operator_raw_thinking": s.get("operator_raw_thinking"), + "operator_native_thinking": s.get("operator_native_thinking"), + "last_execution_result": s.get("last_execution_result"), + "interleaved_events": s.get("interleaved_events"), + } + matched_steps.append(details) + if not matched_steps: + return f"No steps found in range [{s_step}, {e_step}]." + return json.dumps(matched_steps, indent=2, ensure_ascii=False) + + return get_step_details + + +def get_step_screenshot_tool(ctx: ArtemisContext, history_steps: list[dict[str, Any]]) -> BaseTool: + @tool + def get_step_screenshot(step_number: int) -> list[dict[str, Any]] | str: + """Retrieve the screenshot for a specific step. + + Returns a multimodal content block containing the image. + """ + if not ctx.data_engine: + return "Error: DataEngine not available." + + try: + step_num = int(step_number) + except (ValueError, TypeError): + return f"Error: step_number must be an integer, got {step_number}." + + # Find the step to get the image name + step_record = next((s for s in history_steps if s.get("step_number") == step_num), None) + if not step_record: + return f"Error: Step {step_num} not found in history." + + image_name = step_record.get("pre_image_name") or step_record.get("post_image_name") + if not image_name: + return f"Error: No screenshot recorded for step {step_num}." + + images_dir = Path(ctx.data_engine.global_base_dir) / "images" + image_path = images_dir / f"{image_name}.jpg" + + if not image_path.exists(): + return f"Error: Screenshot file for step {step_num} does not exist on disk." + + try: + with open(image_path, "rb") as f: + img_base64 = base64.b64encode(f.read()).decode("utf-8") + except Exception as e: + return f"Error reading screenshot file: {e}" + + # Return a LangChain multimodal content structure + return [ + {"type": "text", "text": f"Screenshot for step {step_num}:"}, + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{img_base64}"}, + }, + ] + + return get_step_screenshot + + +def get_search_history_tool(ctx: ArtemisContext) -> BaseTool: + @tool + def search_history_for_text(query: str) -> str: + """Search for a specific text query within the OCR and UI hierarchy of all historical steps.""" + if not ctx.data_engine: + return "Error: DataEngine not available." + + session_id = ctx.data_engine.current_session_id + if not session_id: + return "Error: No active session." + + try: + steps = ctx.data_engine.storage.get_steps(session_id) + except Exception as e: + return f"Error retrieving steps from storage: {e}" + + results = [] + for step in steps: + step_matches = [] + # Search pre-action image + if step.pre_image_name: + try: + matches = ctx.data_engine.storage.search_ui_by_hash(step.pre_image_name, query) + if matches: + step_matches.extend(matches) + except Exception as e: + logger.error(f"Failed to search UI by hash {step.pre_image_name}: {e}") + + # Search post-action image + if step.post_image_name: + try: + matches = ctx.data_engine.storage.search_ui_by_hash(step.post_image_name, query) + if matches: + step_matches.extend(matches) + except Exception as e: + logger.error(f"Failed to search UI by hash {step.post_image_name}: {e}") + + if step_matches: + # Clean up match details to not bloat context, keep top 3 + clean_matches = [] + for m in step_matches[:3]: + clean_matches.append( + { + "type": m.get("type"), + "matched_text": m.get("matched_text"), + "score": m.get("score"), + } + ) + results.append({"step_number": step.step_number, "matches": clean_matches}) + + if not results: + return f"No matches found for '{query}' in history." + return json.dumps(results, indent=2, ensure_ascii=False) + + return search_history_for_text diff --git a/artemis/agents/planner/planner.json b/artemis/agents/planner/planner.json new file mode 100644 index 0000000..82b4ffa --- /dev/null +++ b/artemis/agents/planner/planner.json @@ -0,0 +1,30 @@ +{ + "blocks": { + "role": "## Role\nYou are the Planner agent. Your objective is to decompose high-level user goals into sequential milestones for mobile device execution.", + "role_validator": "## Role\nYou are the Task Plan Validator. The Operator has attempted to modify the top-level task plan. Your job is to ensure the proposed new plan still aligns with the user's original intent and is viable.", + "guidelines": "## Operating Principles\n\n### 1. Visual Grounding (Start from Current State)\nBase all planning decisions strictly on the provided screenshot of the current device state. Your plan must begin from the current context of the screen.\n\n### 2. Milestone Granularity\nDefine high-level milestones rather than granular UI operations. The milestones you establish are essential checkpoints that the task must pass through; your goal is to keep the task on the right track, not to think about how the task should be completed step-by-step.\n\n### 3. Flexibility and Possibility Retention (Avoid Over-Specification)\nAvoid unnecessarily restricting the execution path or specific tools/applications unless strictly mandated by the user. For example, if the user asks to perform a task that can be accomplished via multiple interfaces (such as a native app or a web browser), describe the milestone in terms of the desired outcome rather than forcing a specific application or interface (e.g. forcing the use of a specific native app when a web browser could also serve). Allow the downstream execution agent the flexibility to choose the most viable path based on the dynamic state of the device.\n\n### 4. Heuristics for Complex & Iterative Operations\nWhen a milestone requires the downstream agent to perform complex iterations, maintain state, or map computed data back to the UI (e.g., adding a pre-selected list of items to a system), you MUST provide heuristics via indented sub-tasks. These sub-tasks should instruct the agent on the exact iterative process (e.g., searching for each item one by one from a note, applying it, and verifying the final state).", + "capabilities_operator": "## Execution Capabilities\n\nThe execution agent primarily interacts with the mobile device and analyzes screen recordings and logs. It can perform the following types of actions:\n- **UI Interaction**: Tapping, swiping, and entering text on the screen.\n- **App Management**: Launching and stopping specific applications.\n- **Information Retrieval & Analysis**: Using OCR to read text on screen, searching the web for external information, querying system logs, and analyzing screen recordings.\n- **Terminal & Background Operations**: Executing arbitrary ADB shell commands on the device for exploration. It is capable of running long-running ADB commands asynchronously in the background as tasks.\n\n### Capability Boundaries & Planning Rules\n- **Cross-App Operations**: When a task requires interacting with multiple applications, split the operations into distinct milestones for each app. Avoid mixing actions for different apps in a single milestone.", + "output_protocol_initial": "## Task Plan & Scratchpad\nSubmit your initial plan by invoking the `save_note` tool with `key='task_plan'`. Avoid writing the plan in your text response, as it will only be recognized if submitted via the tool.\n\nFormat the plan as a Markdown task list starting with `- [ ]` for uncompleted items. You can use indented sub-tasks (e.g., ` - [ ]`) to break down complex exploratory steps and provide heuristics to guide the downstream agent. Ensure the content of the `task_plan` note contains only the task list, excluding any introductions, explanations, or headers.\n\n### Creating Additional Notes\nYou are highly recommended to create additional empty notes (like answer sheets) for the downstream executor to fill out with findings, lists, or analytical information relevant to the task. For example, if the user asks you to analyze specific YouTube videos, you should pre-create an empty note with a descriptive name (e.g., `analysis_of_xxx_videos`), and make sure to put the note name directly inside your milestones to guide the downstream executor in filling out this answer sheet.", + "examples_initial": "## Examples\n\n### Example 1: Basic Task\n**Goal**: \"Send 'I will be late' to Alice on WhatsApp\"\n\n**Incorrect (Too granular)**:\n- [ ] Tap search icon\n- [ ] Type Alice\n- [ ] Tap Alice's avatar\n\n**Correct (Milestones)**:\n- [ ] Open the chat window with Alice in WhatsApp\n- [ ] Successfully send the message \"I will be late\" in Alice's chat window\n\n### Example 2: Information Retrieval & Scratchpad\n**Goal**: \"Find 10 nature videos on YouTube\"\n\n**Incorrect (Vague or directly jumping to end)**:\n- [ ] Open YouTube\n- [ ] Search nature videos and pick 10\n\n**Correct (Leveraging Scratchpad & Sub-tasks Heuristics)**:\n*(Note: As the Planner, pre-create the required empty notes `nature_videos_candidates` and `nature_videos` using the note-taking tools before completing the plan)*\n- [ ] Search for \"nature videos\" on YouTube\n- [ ] Save a list of 15 candidate videos into a temporary note called `nature_videos_candidates`\n - [ ] Add candidate 1\n - [ ] Add candidate 2\n - [ ] (Continue until 15 candidates are found)\n- [ ] Evaluate candidate videos from `nature_videos_candidates` one by one, select exactly 10 correct ones, and save them into `nature_videos`\n - [ ] Evaluate candidate 1: watch, filter out if irrelevant, save if correct\n - [ ] Evaluate candidate 2\n - [ ] (Continue evaluating one by one until 10 are selected)", + "validator_guidelines": "## Validation Principles\n- **Goal Alignment**: Ensure the new plan ultimately accomplishes the `Initial Goal` requested by the user. If the Operator has abandoned the core objective or fundamentally altered it, you must reject the plan.\n- **Logical Coherence**: Verify that the proposed milestones form a logical, achievable sequence to reach the goal. Reject the plan if it introduces contradictory steps, infinite loops, or obvious logical fallacies.\n- **Adaptive Acceptance**: The Operator operates in a dynamic environment and may encounter unexpected UI states. You should approve alternative strategies, extra intermediate steps, or necessary workarounds as long as they serve the primary goal effectively.\n- **Constructive Feedback**: If you must reject the plan, provide clear, objective, and actionable feedback explaining the logical flaw or goal deviation, and suggest how the Operator should correct the plan." + }, + "modes": { + "initial_plan": { + "system": [ + "role", + "guidelines", + "capabilities_operator", + "output_protocol_initial", + "examples_initial" + ], + "human": "Goal: {{ initial_goal }}" + }, + "validator": { + "system": [ + "role_validator", + "validator_guidelines" + ], + "human": "Initial Goal: {{ initial_goal }}\n\n### Execution History\n{{ history_str }}\n\n### Operator's Thoughts Before Modifying\n{% if operator_native_thinking %}Native Thinking:\n{{ operator_native_thinking }}\n\n{% endif %}{% if operator_raw_thinking %}Raw Thinking:\n{{ operator_raw_thinking }}\n\n{% endif %}### Old Plan\n{{ content_before }}\n\n### Proposed New Plan\n{{ content_after }}\n\nPlease evaluate if the proposed new plan should be approved." + } + } +} \ No newline at end of file diff --git a/artemis/agents/planner/planner.py b/artemis/agents/planner/planner.py new file mode 100644 index 0000000..d813c3b --- /dev/null +++ b/artemis/agents/planner/planner.py @@ -0,0 +1,422 @@ +# 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 jinja2 import Template +from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage +from langchain_core.tools import BaseTool +from pydantic import BaseModel, Field + +from artemis.context import ArtemisContext +from artemis.controllers.unified_controller import UnifiedMobileController +from artemis.data_engine.trace import CURRENT_TRACE_ID, TraceSpan, trace +from artemis.graph.state import State +from artemis.services.llm import ( + get_llm, + invoke_llm_with_timeout_message, + with_fallback, +) +from artemis.tools.index import get_tool_by_name, get_tools_from_wrappers +from artemis.tools.scratchpad import ( + list_notes_wrapper, + read_note_wrapper, + save_note_wrapper, + update_note_wrapper, +) +from artemis.tools.tool_wrapper import ( + get_tool_result_content, + invoke_tool_with_injection, +) +from artemis.utils.decorators import wrap_with_callbacks +from artemis.utils.logger import get_logger +from artemis.utils.notes import get_note_file_path +from artemis.utils.task_tree import build_plan_and_history + +logger = get_logger(__name__) + + +class _CyFunctionDetectorMeta(type): + def __instancecheck__(self, instance): + name = type(instance).__name__ + return ( + name + in ( + "cyfunction", + "cython_function_or_method", + "builtin_function_or_method", + ) + or "cyfunction" in name.lower() + ) + + +class CyFunctionDetector(metaclass=_CyFunctionDetectorMeta): + pass + + +class ValidationResult(BaseModel): + model_config = {"ignored_types": (CyFunctionDetector,)} + is_approved: bool = Field( + description=( + "True if the new plan is acceptable and aligns with the initial goal, False otherwise." + ) + ) + feedback: str = Field( + description=( + "If rejected, provide constructive feedback explaining why and what" + " to do instead. If approved, can be empty." + ) + ) + + +async def run_async_planner_validation( + ctx: ArtemisContext, + initial_goal: str, + content_before: str, + content_after: str, + operator_raw_thinking: str | None = None, + operator_native_thinking: str | None = None, +) -> dict: + """Asynchronously validates a top-level plan change proposed by the Operator.""" + logger.info("Starting Async Planner Validation...") + try: + prompts_path = Path(__file__).parent / "planner.json" + with open(prompts_path, encoding="utf-8") as f: + prompts_data = json.load(f) + + mode_config = prompts_data["modes"]["validator"] + system_blocks = mode_config["system"] + system_content = "\n\n".join(prompts_data["blocks"][block] for block in system_blocks) + system_message = Template(system_content).render() + + history_str = "No history available." + if ctx.data_engine: + try: + steps = ctx.data_engine.get_agent_friendly_steps() + history_str = build_plan_and_history( + "", # Pass empty string to avoid rendering the task plan in the history + steps, + "default", + last_n_detailed=2, + strict_milestone_pruning=True, + recent_window_size=5, + chronological_last_step=True, + ) + + except Exception as e: + logger.error(f"Failed to fetch execution history for validator: {e}") + + human_content = mode_config["human"] + human_message = Template(human_content).render( + initial_goal=initial_goal, + content_before=content_before, + content_after=content_after, + operator_raw_thinking=operator_raw_thinking, + operator_native_thinking=operator_native_thinking, + history_str=history_str, + ) + + messages = [ + SystemMessage(content=system_message), + HumanMessage(content=human_message), + ] + + llm = get_llm(ctx=ctx, name="planner").with_structured_output(ValidationResult) + + result: ValidationResult = await invoke_llm_with_timeout_message(llm.ainvoke(messages)) + + return { + "status": "success" if result.is_approved else "failed", + "feedback": result.feedback, + } + except Exception as e: + logger.error(f"Planner async validation failed: {e}") + # In case of error, we fallback to approving so we don't block the agent unnecessarily + return { + "status": "success", + "feedback": f"Validation failed with error: {e}", + } + + +class PlannerNode: + def __init__(self, ctx: ArtemisContext, tools: list[BaseTool] | None = None): + self.ctx = ctx + self.tools = tools or [] + + @wrap_with_callbacks( + before=lambda: logger.info("Starting Planner Agent..."), + on_success=lambda _: logger.success("Planner Agent"), + on_failure=lambda _: logger.error("Planner Agent"), + ) + @trace(type="agent", name="planner") + async def __call__(self, state: State): + prompts_path = Path(__file__).parent / "planner.json" + with open(prompts_path, encoding="utf-8") as f: + prompts_data = json.load(f) + + mode = "initial_plan" + mode_config = prompts_data["modes"][mode] + + system_blocks = mode_config["system"] + system_content = "\n\n".join(prompts_data["blocks"][block] for block in system_blocks) + + system_message = Template(system_content).render() + + # Get screenshot + screenshot_b64 = None + screenshot_path = state.latest_screenshot if hasattr(state, "latest_screenshot") else None + + if screenshot_path and Path(screenshot_path).exists(): + try: + with open(screenshot_path, "rb") as f: + screenshot_b64 = base64.b64encode(f.read()).decode("utf-8") + logger.info(f"Planner loaded screenshot from file: {screenshot_path}") + except Exception as e: + logger.error(f"Failed to read screenshot from {screenshot_path}: {e}") + + # Fallback for Initial Plan: if no screenshot in state, take one on the spot + if not screenshot_b64: + logger.info("No screenshot found in state. Taking screenshot for Initial Plan...") + try: + controller = UnifiedMobileController(self.ctx) + screenshot_b64 = await controller.take_screenshot() + + # Save to Data Engine if available to keep trace + if self.ctx.data_engine: + img_bytes = base64.b64decode(screenshot_b64) + image_name = self.ctx.data_engine.get_or_create_image(img_bytes) + saved_path = str(self.ctx.data_engine.get_image_path(image_name)) + # Update local state object so it can be used later if needed + state.latest_screenshot = saved_path + logger.info(f"Saved initial plan screenshot to {saved_path}") + except Exception as e: + logger.error(f"Failed to take screenshot for Initial Plan: {e}") + + if not screenshot_b64: + logger.error("Planner failed to obtain a screenshot.") + raise ValueError("Planner failed to obtain a screenshot.") + + human_content = mode_config["human"] + render_kwargs = { + "initial_goal": state.initial_goal, + } + + human_message = Template(human_content).render(**render_kwargs) + + human_message_content = [{"type": "text", "text": human_message}] + if screenshot_b64: + human_message_content.append( + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{screenshot_b64}"}, + } + ) + messages = [ + SystemMessage(content=system_message), + HumanMessage(content=human_message_content), + ] + + llm = get_llm(ctx=self.ctx, name="planner") + llm_fallback = get_llm(ctx=self.ctx, name="planner", use_fallback=True) + + # Ensure tools are populated (fallback if not passed in __init__, e.g. in tests) + if not self.tools: + note_wrappers = [ + save_note_wrapper, + read_note_wrapper, + update_note_wrapper, + list_notes_wrapper, + ] + self.tools = get_tools_from_wrappers(self.ctx, note_wrappers) + + all_tools = list(self.tools) + + if all_tools: + llm = llm.bind_tools(all_tools) + llm_fallback = llm_fallback.bind_tools(all_tools) + + def validate_plan_format(content: str) -> tuple[bool, str]: + lines = content.splitlines() + has_subgoal = False + for line in lines: + stripped = line.strip() + if stripped.startswith("- ["): + has_subgoal = True + if not ( + stripped.startswith("- [ ]") + or stripped.startswith("- [x]") + or stripped.startswith("- [!]") + ): + return ( + False, + ( + f"Invalid task status in line: {line}. Must be" + " '[ ]', '[x]', or '[!]'." + ), + ) + elif stripped.startswith("#") or not stripped: + continue + else: + continue + + if not has_subgoal: + return ( + False, + ("Plan must contain at least one subgoal starting with '- [ ]'."), + ) + + return True, "" + + current_messages = messages + max_validation_attempts = 3 + max_iterations = 10 + + for validation_attempt in range(max_validation_attempts): + tool_called = False + + async def run_stream(llm_model): + full_response = None + trace_id = CURRENT_TRACE_ID.get() + async for chunk in llm_model.astream(current_messages): + if full_response is None: + full_response = chunk + else: + full_response += chunk + + pass + return full_response + + for _ in range(max_iterations): + response = await with_fallback( + main_call=lambda: invoke_llm_with_timeout_message(run_stream(llm)), + fallback_call=lambda: invoke_llm_with_timeout_message(run_stream(llm_fallback)), + ) + + if response is None or not response.tool_calls: + break + + current_messages.append(response) + + for tc in response.tool_calls: + tool_name = tc["name"] + logger.info(f"Planner called tool: {tool_name} with args: {tc['args']}") + if ":" in tool_name: + tool_to_run = get_tool_by_name(tool_name, all_tools) + else: + tool_to_run = next((t for t in all_tools if t.name == tool_name), None) + + if tool_to_run and tool_to_run.name in [ + "save_note", + "update_note", + ]: + tool_called = True + if tool_to_run: + args = dict(tc["args"]) + with TraceSpan(name=tool_name, ctx=self.ctx) as span: + span.payload = {"args": args} + try: + logger.info( + f"Tool {tool_name} expected args:" + f" {list(tool_to_run.args.keys())}" + ) + + logger.info(f"Invoking {tool_name} with args: {list(args.keys())}") + result_obj = await invoke_tool_with_injection( + tool=tool_to_run, + args=args, + tool_call_id=tc["id"], + state=state, + ) + result_content = get_tool_result_content(result_obj) + span.result = result_content + + current_messages.append( + ToolMessage( + tool_call_id=tc["id"], + content=result_content, + ) + ) + except Exception as e: + span.status = "failed" + span.error = str(e) + current_messages.append( + ToolMessage( + tool_call_id=tc["id"], + content=f"Error: {e}", + status="error", + ) + ) + else: + current_messages.append( + ToolMessage( + tool_call_id=tc["id"], + content=f"Error: Tool {tool_name} not found", + status="error", + ) + ) + + # Post-execution validation + plan_exists = False + plan_valid = False + error_message = "" + + if not tool_called: + error_message = ( + "You did not call `save_note` or `update_note` to submit/update the plan." + ) + elif self.ctx.data_engine: + task_plan_path = get_note_file_path(self.ctx.data_engine.base_dir, "task_plan") + if task_plan_path.exists(): + plan_exists = True + try: + content = task_plan_path.read_text(encoding="utf-8") + plan_valid, error_message = validate_plan_format(content) + except Exception as e: + error_message = f"Failed to read task plan for validation: {e}" + else: + error_message = ( + "Task plan under key 'task_plan' does not exist despite tool call." + ) + + if tool_called and plan_exists and plan_valid: + logger.info("Plan submitted and validated successfully.") + break + + logger.warning( + f"Plan validation failed (attempt {validation_attempt + 1}): {error_message}" + ) + + # Feedback to model + current_messages.append( + HumanMessage( + content=( + f"Plan validation failed: {error_message}\nPlease" + " correct the plan and try again by calling the" + " appropriate tool." + ) + ) + ) + + else: + logger.error( + f"Failed to generate a valid plan after {max_validation_attempts} attempts." + ) + + return await state.asanitize_update( + ctx=self.ctx, + update={}, + agent="planner", + ) diff --git a/artemis/agents/planner/planner_agent.py b/artemis/agents/planner/planner_agent.py new file mode 100644 index 0000000..8fb6946 --- /dev/null +++ b/artemis/agents/planner/planner_agent.py @@ -0,0 +1,53 @@ +# 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. + +"""Planner Agent implementation adhering to BaseAgent lifecycle.""" + +from typing import Any +from artemis.agents.base import AgentConfig, AgentRegistry, BaseAgent +from artemis.agents.planner.planner import PlannerNode +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +@AgentRegistry.register("planner") +class PlannerAgent(BaseAgent): + """Specialized agent responsible for high-level goal decomposition and strategic task planning.""" + + def __init__( + self, + config: AgentConfig | None = None, + ctx: Any | None = None, + **kwargs: Any, + ): + super().__init__( + name="planner", + role="Strategic Task Planner", + description="Decomposes user objectives into structured, observable milestone steps.", + config=config, + ctx=ctx, + ) + self._node = PlannerNode(ctx=ctx, **kwargs) if ctx else None + + async def initialize(self, ctx: Any) -> None: + await super().initialize(ctx) + self._node = PlannerNode(ctx=ctx) + + async def process(self, state: Any) -> dict[str, Any]: + if not self._node and self.ctx: + self._node = PlannerNode(ctx=self.ctx) + if not self._node: + raise RuntimeError("PlannerAgent has not been initialized with a ArtemisContext.") + return await self._node(state) diff --git a/artemis/agents/planner/types.py b/artemis/agents/planner/types.py new file mode 100644 index 0000000..8cc1b5d --- /dev/null +++ b/artemis/agents/planner/types.py @@ -0,0 +1,85 @@ +# 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. + +from datetime import datetime +from enum import Enum +from typing import Annotated + +from pydantic import BaseModel + + +class _CyFunctionDetectorMeta(type): + def __instancecheck__(self, instance): + name = type(instance).__name__ + return ( + name + in ( + "cyfunction", + "cython_function_or_method", + "builtin_function_or_method", + ) + or "cyfunction" in name.lower() + ) + + +class CyFunctionDetector(metaclass=_CyFunctionDetectorMeta): + pass + + +class PlannerSubgoalOutput(BaseModel): + model_config = {"ignored_types": (CyFunctionDetector,)} + description: str + + +class PlannerOutput(BaseModel): + model_config = {"ignored_types": (CyFunctionDetector,)} + subgoals: list[PlannerSubgoalOutput] + + +class SubgoalStatus(Enum): + NOT_STARTED = "NOT_STARTED" + PENDING = "PENDING" + SUCCESS = "SUCCESS" + FAILURE = "FAILURE" + + +class Subgoal(BaseModel): + model_config = {"ignored_types": (CyFunctionDetector,)} + id: Annotated[str, "Unique identifier of the subgoal"] + description: Annotated[str, "Description of the subgoal"] + completion_reason: Annotated[ + str | None, "Reason why the subgoal was completed (failure or success)" + ] = None + status: SubgoalStatus + started_at: Annotated[datetime | None, "When the subgoal started"] = None + ended_at: Annotated[datetime | None, "When the subgoal ended"] = None + + def __str__(self): + status_emoji = "โ“" + if self.status == SubgoalStatus.SUCCESS: + status_emoji = "โœ…" + elif self.status == SubgoalStatus.FAILURE: + status_emoji = "โŒ" + elif self.status == SubgoalStatus.PENDING: + status_emoji = "โณ" + elif self.status == SubgoalStatus.NOT_STARTED: + status_emoji = "(not started yet)" + + output = f"- [ID:{self.id}]: {self.description} : {status_emoji}." + if self.completion_reason: + output += f" Completion reason: {self.completion_reason}" + return output + + def __repr__(self): + return str(self) diff --git a/artemis/agents/planner/utils.py b/artemis/agents/planner/utils.py new file mode 100644 index 0000000..2d8a0e8 --- /dev/null +++ b/artemis/agents/planner/utils.py @@ -0,0 +1,89 @@ +# 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. + +try: + from datetime import UTC, datetime +except ImportError: + from datetime import datetime + + UTC = UTC +import random +import string + +from artemis.agents.planner.types import Subgoal, SubgoalStatus + + +def get_current_subgoal(subgoals: list[Subgoal]) -> Subgoal | None: + return next((s for s in subgoals if s.status == SubgoalStatus.PENDING), None) + + +def get_subgoals_by_ids(subgoals: list[Subgoal], ids: list[str]) -> list[Subgoal]: + return [s for s in subgoals if s.id in ids] + + +def get_next_subgoal(subgoals: list[Subgoal]) -> Subgoal | None: + return next((s for s in subgoals if s.status == SubgoalStatus.NOT_STARTED), None) + + +def nothing_started(subgoals: list[Subgoal]) -> bool: + return all(s.status == SubgoalStatus.NOT_STARTED for s in subgoals) + + +def complete_current_subgoal(subgoals: list[Subgoal]) -> list[Subgoal]: + current_subgoal = get_current_subgoal(subgoals) + if not current_subgoal: + return subgoals + current_subgoal.status = SubgoalStatus.SUCCESS + current_subgoal.ended_at = datetime.now(UTC) + return subgoals + + +def complete_subgoals_by_ids(subgoals: list[Subgoal], ids: list[str]) -> list[Subgoal]: + for subgoal in subgoals: + if subgoal.id in ids: + subgoal.status = SubgoalStatus.SUCCESS + subgoal.ended_at = datetime.now(UTC) + return subgoals + + +def fail_current_subgoal(subgoals: list[Subgoal]) -> list[Subgoal]: + current_subgoal = get_current_subgoal(subgoals) + if not current_subgoal: + return subgoals + current_subgoal.status = SubgoalStatus.FAILURE + current_subgoal.ended_at = datetime.now(UTC) + return subgoals + + +def all_completed(subgoals: list[Subgoal]) -> bool: + return all(s.status == SubgoalStatus.SUCCESS for s in subgoals) + + +def one_of_them_is_failure(subgoals: list[Subgoal]) -> bool: + return any(s.status == SubgoalStatus.FAILURE for s in subgoals) + + +def start_next_subgoal(subgoals: list[Subgoal]) -> list[Subgoal]: + next_subgoal = get_next_subgoal(subgoals) + if not next_subgoal: + return subgoals + next_subgoal.status = SubgoalStatus.PENDING + next_subgoal.started_at = datetime.now(UTC) + return subgoals + + +def generate_id(length: int = 6) -> str: + """Generates a small and distinct random string ID.""" + chars = string.ascii_lowercase + string.digits + return "".join(random.choice(chars) for _ in range(length)) diff --git a/artemis/agents/summarizer/summarizer.md b/artemis/agents/summarizer/summarizer.md new file mode 100644 index 0000000..a3612c3 --- /dev/null +++ b/artemis/agents/summarizer/summarizer.md @@ -0,0 +1,31 @@ +# ROLE +You are the Summarizer Agent for a mobile automation system. Your task is to summarize the latest execution step, providing a high-information-density and succinct history summary of the current step without losing any information. + +# OUTPUT FORMAT +- **Format Constraint**: You MUST output the summary as a **single, continuous paragraph**. +- **Prohibited Formats**: Do NOT use bullet points, numbered lists, markdown formatting, or multiple paragraphs. +- **Goal**: Maintain a consistent, easily parsable narrative for the Operator's historical record. + +# CORE PHILOSOPHY & RULES + +1. **Perspective**: You MUST write the summary from the Operator's first-person perspective using "I" (e.g., "I intended to...", "I clicked..."). NEVER use third-person terms like "The operator" or "The agent", nor second-person terms like "You" when referring to the Operator. For other agents (e.g., Failure Analyzer, Explorer), you can use their names. +2. **Zero Subjective Validation**: Absolutely avoid declaring semantic success, completion, or failure. In historical compression, any subjective validation will cause future planners to falsely believe a goal was definitively met or permanently broken. + - **BANNED WORDS**: "successfully", "completed", "entered", "navigated to", "failed", "unsuccessful", "could not", "achieved". +3. **Focus on Intent & Strategy**: The specific physical action (exact coordinates, swipes) is now recorded automatically. Focus primarily on the high-level Intent and reasoning. If you realize the Operator is on the wrong path and needs to backtrack or pivot, describe this as a shift in strategy. +4. **Intended Context Transitions**: Describe the intended navigation path ("From Where -> Intended Where"). Do not declare that the transition actually happened. (e.g., "From the Home screen, I intended to open the Network Settings view"). +5. **Preserve Verifications & Missing Items**: If the Operator verified any state in this step, or explicitly mentioned what prerequisites or items are still missing to accomplish the goal, you MUST state these contents in detail in the summary. +6. **Objective Intervention Logs**: If the Failure Analyzer intervened, faithfully summarize what triggered the intervention and what the analyzer executed, without judging the outcome. + +# LOOP & STAGNATION DETECTION +Use the `RECENT 10 STEPS HISTORY` to intelligently detect loops or stagnation: +1. **Physical Action Loops**: Be vigilant for loops spanning multiple steps based on physical interactions. If you detect a valid action loop, explicitly mention the pattern and the current iteration count (e.g., "This is the 3rd cycle of attempting to access the video group folder"). +2. **Cognitive Stagnation**: While you generally ignore internal background tool calls (e.g., reading logs, checking states) when counting physical action loops, there is an exception: if you detect continuous multiple steps of purely internal tool calls without any physical actions, you MUST objectively record this strategic "stagnation" or "analysis loop". + +# INPUT STRUCTURE +You will receive: +- **RECENT 10 STEPS HISTORY**: A brief history of the last 10 steps to help you detect loops and stagnation patterns. +- **CHRONOLOGICAL STEP TRACE (CURRENT STEP)**: The step-by-step chronological history of the current execution turn, including: + - `[Operator Monologue] / [Operator Native Thought]`: The operator's planning, intentions, progress counting, reasoning, verified states, and explicit mentions of missing prerequisites. + - `[Operator Tool Call] / [Operator Final Action]`: The actual action(s) planned and executed. + - `[Pre-Execution Safety Net]`: Validation checks prior to action execution. + - `[Failure Analyzer Recovery Loop]`: Interventions, thoughts, auxiliary tool calls, and outcomes of the Failure Analyzer. diff --git a/artemis/agents/summarizer/summarizer.py b/artemis/agents/summarizer/summarizer.py new file mode 100644 index 0000000..42455ba --- /dev/null +++ b/artemis/agents/summarizer/summarizer.py @@ -0,0 +1,231 @@ +# 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 asyncio +import json +from pathlib import Path +from uuid import UUID + +from langchain_core.messages import ( + HumanMessage, + SystemMessage, +) + +from artemis.context import ArtemisContext +from artemis.data_engine.trace import CURRENT_TRACE_ID, trace +from artemis.graph.state import State +from artemis.services.llm import get_llm +from artemis.utils.coordinates import normalize_any_structure +from artemis.utils.logger import get_logger +from artemis.utils.task_tree import ( + build_plan_and_history, + get_active_subgoal_hashes, +) + +logger = get_logger(__name__) + + +class SummarizerNode: + def __init__(self, ctx: ArtemisContext): + self.ctx = ctx + + @trace(type="agent", name="summarizer") + async def __call__(self, state: State): + # 1. Trigger background task for summary and data engine update + if self.ctx.data_engine and state.current_step_id: + step_id = UUID(state.current_step_id) + decisions = state.structured_decisions + operator_raw_thinking = getattr(state, "operator_raw_thinking", None) + operator_native_thinking = getattr(state, "operator_native_thinking", None) + last_execution_result = getattr(state, "last_execution_result", None) + + # Spawn background task + trace_id = CURRENT_TRACE_ID.get() + task = asyncio.create_task( + self._generate_summary_and_update_engine( + step_id, + decisions, + operator_raw_thinking, + operator_native_thinking, + last_execution_result, + trace_id, + ), + name=f"summarizer_step_{step_id}", + ) + if hasattr(self.ctx, "background_tasks") and isinstance( + self.ctx.background_tasks, list + ): + self.ctx.background_tasks.append(task) + + return {} + + async def _generate_summary_and_update_engine( + self, + step_id, + decisions, + operator_raw_thinking, + operator_native_thinking, + last_execution_result, + trace_id, + ): + """Background task to generate summary and update Data Engine with retries.""" + + max_retries = 3 + retry_delay = 2 + + for attempt in range(max_retries): + try: + logger.info(f"Generating summary for step {step_id} (attempt {attempt + 1})") + + # Prepare prompt for summary + llm = get_llm(ctx=self.ctx, name="summarizer") + + prompt_path = Path(__file__).parent.joinpath("summarizer.md") + if not prompt_path.exists(): + prompt = ( + "You are a summarizer. Summarize the latest action and" + " its result in 1-2 concise sentences. Focus on: Action" + " taken -> Result -> Observation." + ) + else: + prompt = prompt_path.read_text(encoding="utf-8") + + # Normalize coordinates before presenting to the summarizer LLM + width = ( + getattr(self.ctx.device, "device_width", 1080) + if self.ctx and self.ctx.device + else 1080 + ) + height = ( + getattr(self.ctx.device, "device_height", 2400) + if self.ctx and self.ctx.device + else 2400 + ) + + content = [] + + # Consistently use optimal build_plan_and_history to construct accurate context + plan_and_history_str = "" + if self.ctx.data_engine: + try: + friendly_steps = self.ctx.data_engine.get_agent_friendly_steps() + if isinstance(friendly_steps, list): + # 1. Trim history to only include steps up to the current step being summarized (inclusive) + current_steps = [] + for s in friendly_steps: + current_steps.append(s) + if str(s.get("step_id")) == str(step_id): + break + + # 2. Read task plan to provide goal orientation to the summarizer + task_plan = "" + base_dir = getattr(self.ctx.data_engine, "base_dir", None) + if base_dir and isinstance(base_dir, (str, Path)): + notes_dir = Path(base_dir) / "notes" + task_plan_path = notes_dir / "task_plan.md" + if task_plan_path.exists(): + task_plan = task_plan_path.read_text(encoding="utf-8") + + # 3. Unified rendering: concise for first N steps (with self-healing facts), detailed for current step (without redundant summary header) + + subgoal_hash = "default" + if task_plan: + subgoal_hash, _ = get_active_subgoal_hashes(task_plan) + + plan_and_history_str = build_plan_and_history( + task_plan=task_plan, + steps=current_steps, + current_subgoal_hash=subgoal_hash, + last_n_detailed=1, + min_summaries=10, + strict_milestone_pruning=False, + ) + except Exception as e: + logger.error(f"Failed to build plan and history for Summarizer: {e}") + + if plan_and_history_str: + content.append(plan_and_history_str) + else: + # Fallback to standard raw fields if chronological trace couldn't be generated + if decisions: + try: + dec_obj = json.loads(decisions) + dec_obj = normalize_any_structure(dec_obj, width, height) + content.append( + f"Decisions made: {json.dumps(dec_obj, ensure_ascii=False)}" + ) + except Exception: + content.append(f"Decisions made: {decisions}") + + if operator_raw_thinking: + content.append(f"Operator explicit thoughts:\n{operator_raw_thinking}") + + if operator_native_thinking: + content.append(f"Operator native thoughts:\n{operator_native_thinking}") + + messages = [ + SystemMessage(content=prompt), + HumanMessage(content="\n\n".join(content)), + ] + + generated_text = "" + aggregated_chunk = None + async for chunk in llm.astream(messages): + if aggregated_chunk is None: + aggregated_chunk = chunk + else: + aggregated_chunk += chunk + + if chunk.content: + text_to_stream = "" + if isinstance(chunk.content, str): + text_to_stream = chunk.content + elif isinstance(chunk.content, list): + for item in chunk.content: + if isinstance(item, dict) and item.get("type") == "text": + text_to_stream += item.get("text", "") + + if text_to_stream: + generated_text += text_to_stream + pass + + summary = generated_text.strip() + logger.info(f"Summary generated: {summary}") + + if aggregated_chunk and getattr(aggregated_chunk, "usage_metadata", None): + usage = aggregated_chunk.usage_metadata + cached_tokens = usage.get("input_token_details", {}).get("cache_read", 0) + logger.info( + "Summarizer LLM usage:" + f" prompt_tokens={usage.get('input_tokens')}" + f" (cached={cached_tokens})," + f" completion_tokens={usage.get('output_tokens')}" + ) + + # Update Data Engine + self.ctx.data_engine.update_step_summary(step_id, summary) + logger.info(f"Data Engine updated with summary for step {step_id}") + return # Success + + except Exception as e: + logger.error( + f"Failed to generate summary or update engine (attempt {attempt + 1}): {e}" + ) + if attempt < max_retries - 1: + await asyncio.sleep(retry_delay) + retry_delay *= 2 # Exponential backoff + else: + logger.error( + f"Max retries reached for step {step_id}. Summary generation failed." + ) diff --git a/artemis/agents/summarizer/summarizer_agent.py b/artemis/agents/summarizer/summarizer_agent.py new file mode 100644 index 0000000..1331d46 --- /dev/null +++ b/artemis/agents/summarizer/summarizer_agent.py @@ -0,0 +1,53 @@ +# 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. + +"""Summarizer Agent implementation adhering to BaseAgent lifecycle.""" + +from typing import Any +from artemis.agents.base import AgentConfig, AgentRegistry, BaseAgent +from artemis.agents.summarizer.summarizer import SummarizerNode +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +@AgentRegistry.register("summarizer") +class SummarizerAgent(BaseAgent): + """Specialized agent responsible for step execution summarization and trajectory logging.""" + + def __init__( + self, + config: AgentConfig | None = None, + ctx: Any | None = None, + **kwargs: Any, + ): + super().__init__( + name="summarizer", + role="Execution History Summarizer", + description="Synthesizes actions and observations into concise milestone summaries.", + config=config, + ctx=ctx, + ) + self._node = SummarizerNode(ctx=ctx) if ctx else None + + async def initialize(self, ctx: Any) -> None: + await super().initialize(ctx) + self._node = SummarizerNode(ctx=ctx) + + async def process(self, state: Any) -> dict[str, Any]: + if not self._node and self.ctx: + self._node = SummarizerNode(ctx=self.ctx) + if not self._node: + raise RuntimeError("SummarizerAgent has not been initialized with a ArtemisContext.") + return await self._node(state) diff --git a/artemis/agents/validator/failure_analyzer.md b/artemis/agents/validator/failure_analyzer.md new file mode 100644 index 0000000..9d7e5b4 --- /dev/null +++ b/artemis/agents/validator/failure_analyzer.md @@ -0,0 +1,25 @@ +# Identity & Goal +You are a specialized error recovery agent. You will be provided with a failed action. Your ONLY task is to repair this specific failure and restore the system state so the main system can resume. + +# Context +- **Initial Goal**: {{ initial_goal }} +- **History & Plan**: +{{ plan_and_history }} + +# Core Workflow +1. **Analyze**: Compare the two states provided to you: + - **Decision State (`Screenshot Seen During System Decision`)**: This is the screen the system saw when it originally planned the action. It shows the target element successfully located. + - **Failed State (`Latest Screenshot (Failed State)`)**: This is the live screen captured at the moment you were summoned. It shows the current state where the target is missing or unresponsive. + Your goal in this comparison is to: + - Read the Operator's thought process in the decision loop of the failed step to understand what they were originally trying to accomplish with this failed action. + - Understand the expected target layout. + - Identify the cause of the failure based on the error message and the visual mismatch. +2. **Recover & Execute**: Restore the target state and execute the failed action. + - **Flexibility**: Focus on outcome over form. Take alternative/adaptive steps (e.g., scroll, navigate, different click paths) if original coordinates are invalid. +3. **Report**: Call `report_failure_analysis` to conclude: + - `status="fixed"`: Successfully executed the failed action, reached the destination state, or false alarm. + - `status="cannot_fix"`: Target is unrecoverable, the app is broken, or the 15-step budget is exceeded. + - `analysis`: Summarize the root cause, the fix applied, and explicitly confirming that control is handed back to the main system. + +# Constraints +The most recent step (Step {{ failed_step_number }}: {{ failed_action_description }}) failed to execute. You must resolve the issue and report the results in 15 interactive steps or fewer. Now please begin the repair analysis. Please note that during the **Recover & Execute** phase, you only need to resolve the currently failed step; do not attempt to execute any subsequent steps or complete the overall goal yourself. Once the failure is resolved, please immediately enter the **Report** phase and return control to the system. diff --git a/artemis/agents/validator/failure_analyzer.py b/artemis/agents/validator/failure_analyzer.py new file mode 100644 index 0000000..735a0ea --- /dev/null +++ b/artemis/agents/validator/failure_analyzer.py @@ -0,0 +1,607 @@ +# 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. + +"""Universal Multi-Model Failure Analyzer for Artemis. + +Analyzes execution failures and performs autonomous self-healing across +Google Gemini, OpenAI GPT-4o/o3, Anthropic Claude 3.5/3.7, and OpenRouter endpoints. +""" + +import asyncio +from enum import StrEnum +import json +from pathlib import Path +from typing import Any +import uuid + +from jinja2 import Template +from langchain_core.messages import ( + BaseMessage, + HumanMessage, + SystemMessage, + ToolMessage, +) + +from artemis.agents.validator.tool_declarations import ( + ASK_EXPLORER_TOOL, + CLICK_SEQUENCE_TOOL, + MobileActionExecutor, + ToolDeclaration, + VALIDATOR_TOOLS_DECLARATION, + capture_screenshot_and_parse_ui, + prune_intermediate_screenshots, +) +from artemis.context import ArtemisContext +from artemis.controllers.unified_controller import UnifiedMobileController +from artemis.data_engine.trace import trace +from artemis.graph.state import State +from artemis.services.llm import get_llm +from artemis.utils.logger import get_logger +from artemis.utils.notes import get_note_file_path +from artemis.utils.ocr_xml_fusion import fuse_ocr_with_xml +from artemis.utils.task_tree import ( + build_plan_and_history, + format_action_clean, + get_active_subgoal_hashes, +) +from artemis.utils.visualization import format_minimal_list_with_elements + + +class ValidationErrorCategory(StrEnum): + NONE = "none" + TARGET_DISAPPEARED = "target_disappeared" + TARGET_SHIFTED = "target_shifted" + TARGET_OCCUPIED = "target_occupied" + PIXEL_TARGET_DISAPPEARED = "pixel_target_disappeared" + PIXEL_BYPASSED = "pixel_bypassed" + XML_BYPASSED = "xml_bypassed" + GENERAL = "general" + + +logger = get_logger(__name__) + + +# ============================================================================== +# Backward-Compatibility Helper Aliases +# ============================================================================== + + +async def _capture_screenshot_and_parse_ui( + ctx: ArtemisContext, + state: State, + controller: UnifiedMobileController, + skip_settling: bool = False, +) -> tuple[str | None, bytes | None, str | None]: + return await capture_screenshot_and_parse_ui( + ctx, state, controller, skip_settling=skip_settling + ) + + +async def _exec_click( + target: list[int], + state: State, + controller: UnifiedMobileController, + ctx: ArtemisContext, + times: int = 1, + delay_ms: int = 100, +) -> tuple[str, bytes | None, str | None, str | None]: + return await MobileActionExecutor(ctx, controller).exec_click( + target, state, times=times, delay_ms=delay_ms + ) + + +async def _exec_click_sequence( + sequence: list[int | list[int]], + state: State, + controller: UnifiedMobileController, + ctx: ArtemisContext, +) -> tuple[str, bytes | None, str | None, str | None]: + return await MobileActionExecutor(ctx, controller).exec_click_sequence(sequence, state) + + +async def _exec_long_press( + target: list[int], + state: State, + controller: UnifiedMobileController, + ctx: ArtemisContext, + duration_ms: int = 1000, + duration: int | None = None, +) -> tuple[str, bytes | None, str | None, str | None]: + return await MobileActionExecutor(ctx, controller).exec_long_press( + target, state, duration_ms=duration_ms, duration=duration + ) + + +async def _exec_input_text( + text: str, + target: list[int] | None, + state: State, + controller: UnifiedMobileController, + ctx: ArtemisContext, + clear_exist: bool = True, +) -> tuple[str, bytes | None, str | None, str | None]: + return await MobileActionExecutor(ctx, controller).exec_input_text( + text, target, state, clear_exist=clear_exist + ) + + +async def _exec_swipe( + action: str | list[int], + duration: int = 400, + state: State | None = None, + controller: UnifiedMobileController | None = None, + ctx: ArtemisContext | None = None, + duration_ms: int | None = None, +) -> tuple[str, bytes | None, str | None, str | None]: + return await MobileActionExecutor(ctx, controller).exec_swipe( + action, duration=duration, state=state, duration_ms=duration_ms + ) + + +async def _exec_press_key( + key: str, + state: State, + controller: UnifiedMobileController, + ctx: ArtemisContext, +) -> tuple[str, bytes | None, str | None, str | None]: + return await MobileActionExecutor(ctx, controller).exec_press_key(key, state) + + +async def _exec_manage_app( + action: str, + app_name: str, + state: State, + controller: UnifiedMobileController, + ctx: ArtemisContext, +) -> tuple[str, bytes | None, str | None, str | None]: + return await MobileActionExecutor(ctx, controller).exec_manage_app(action, app_name, state) + + +async def _exec_wait_for_delay( + time_in_ms: int, + state: State, + controller: UnifiedMobileController, + ctx: ArtemisContext, +) -> tuple[str, bytes | None, str | None, str | None]: + return await MobileActionExecutor(ctx, controller).exec_wait_for_delay(time_in_ms, state) + + +async def _exec_wait_for_text( + text: str, + wait_state: str | None, + timeout_ms: int | None, + state: State, + controller: UnifiedMobileController, + ctx: ArtemisContext, +) -> tuple[str, bytes | None, str | None, str | None]: + return await MobileActionExecutor(ctx, controller).exec_wait_for_text( + text, wait_state, timeout_ms, state + ) + + +async def _exec_read_note( + key: str, + ctx: ArtemisContext, + start_line: int | None = None, + end_line: int | None = None, +) -> str: + return await MobileActionExecutor(ctx).exec_read_note( + key, start_line=start_line, end_line=end_line + ) + + +async def _exec_list_notes(ctx: ArtemisContext) -> str: + return await MobileActionExecutor(ctx).exec_list_notes() + + +async def _exec_ask_explorer( + query: str, + context_feedback: str | None, + state: State, + ctx: ArtemisContext, + version: str = "flash", +) -> str: + return await MobileActionExecutor(ctx).exec_ask_explorer( + query, context_feedback, state, version=version + ) + + +def _get_ask_explorer_tool_declaration() -> ToolDeclaration: + return ASK_EXPLORER_TOOL + + +class FailureAnalysisStrategy: + """Base strategy class for FailureAnalyzer strategies.""" + + def get_prompt_template_name(self) -> str: + return "failure_analyzer.md" + + def get_tools(self) -> list: + return VALIDATOR_TOOLS_DECLARATION + + def get_user_message_suffix(self, pre_screenshot: bool, post_screenshot: bool) -> str: + suffix = ( + "Your objective is:\n" + "1. Analyze why the action planned in the failed step failed, using the failed step's" + " details, the decision screenshot (Screenshot Seen During System Decision), the" + " latest failed screenshot (Latest Screenshot (Failed State)), and their corresponding" + " UI Element lists.\n" + "2. Resolve any timing, state mismatches (e.g., re-triggering disappeared elements), or" + " environmental obstacles using your execution tools.\n" + "3. Crucially: You MUST execute the failed action yourself using your execution tools" + " to completion.\n" + "4. Once you have successfully executed the repair and the failed action, call the" + " `report_failure_analysis` tool with status='fixed'.\n" + "5. If you are not confident or cannot resolve the failure, call" + " `report_failure_analysis` with status='cannot_fix'." + ) + if pre_screenshot and post_screenshot: + suffix += ( + "\nI have attached the screenshot seen when the system made its" + " decision (Screenshot Seen During System Decision) and the" + " latest screenshot in the failed state (Latest Screenshot" + " (Failed State))." + ) + return suffix + + +class TargetDisappearedStrategy(FailureAnalysisStrategy): + """Specialized strategy for target disappeared errors.""" + + def get_prompt_template_name(self) -> str: + return "target_disappeared_analyzer.md" + + def get_tools(self) -> list: + tools = list(VALIDATOR_TOOLS_DECLARATION) + tools.insert(1, CLICK_SEQUENCE_TOOL) + tools.append(ASK_EXPLORER_TOOL) + return tools + + def get_user_message_suffix(self, pre_screenshot: bool, post_screenshot: bool) -> str: + return "" + + +class PixelTargetDisappearedStrategy(FailureAnalysisStrategy): + """Specialized strategy for pixel-level target disappeared errors.""" + + def get_prompt_template_name(self) -> str: + return "pixel_target_disappeared_analyzer.md" + + def get_tools(self) -> list: + tools = list(VALIDATOR_TOOLS_DECLARATION) + tools.insert(1, CLICK_SEQUENCE_TOOL) + tools.append(ASK_EXPLORER_TOOL) + return tools + + def get_user_message_suffix(self, pre_screenshot: bool, post_screenshot: bool) -> str: + return "" + + +class FailureAnalyzer: + """Universal Multi-Model Failure Analyzer.""" + + def __init__(self, ctx: ArtemisContext): + self.ctx = ctx + self.max_iterations = 15 + + def _select_strategy(self, error_category: ValidationErrorCategory) -> FailureAnalysisStrategy: + if error_category == ValidationErrorCategory.PIXEL_TARGET_DISAPPEARED: + logger.info("Routing failure analysis to specialized PixelTargetDisappearedStrategy.") + return PixelTargetDisappearedStrategy() + if error_category == ValidationErrorCategory.TARGET_DISAPPEARED: + logger.info("Routing failure analysis to specialized TargetDisappearedStrategy.") + return TargetDisappearedStrategy() + logger.info("Routing failure analysis to DefaultAnalysisStrategy.") + return FailureAnalysisStrategy() + + def _get_active_subgoal_hash(self) -> str: + if not self.ctx.data_engine: + return "default" + task_plan_path = get_note_file_path(self.ctx.data_engine.base_dir, "task_plan") + if not task_plan_path.exists(): + return "default" + try: + content = task_plan_path.read_text(encoding="utf-8") + parent_hash, _ = get_active_subgoal_hashes(content) + return parent_hash + except Exception as e: + logger.error(f"Failed to parse active subgoal: {e}") + return "default" + + def _build_plan_and_history(self, steps: list[dict], task_plan: str) -> str: + active_subgoal_hash = self._get_active_subgoal_hash() + return build_plan_and_history( + task_plan, + steps, + active_subgoal_hash, + last_n_detailed=1, + strict_milestone_pruning=True, + recent_window_size=3, + chronological_last_step=True, + for_failure_analyzer=True, + ) + + def _prune_intermediate_screenshots(self, messages: list[Any]): + """Prunes binary screenshot blocks from intermediate messages.""" + prune_intermediate_screenshots(messages) + + @trace(type="agent", name="failure_analyzer") + async def analyze( + self, + state: State, + failed_action: dict, + error_msg: str, + pre_screenshot: str | None = None, + post_screenshot: str | None = None, + pre_screenshot_name: str | None = None, + post_screenshot_name: str | None = None, + executed_actions: list[dict] | None = None, + unexecuted_actions: list[dict] | None = None, + error_category: ValidationErrorCategory = ValidationErrorCategory.GENERAL, + ) -> dict: + if not post_screenshot or post_screenshot == "None": + try: + controller = UnifiedMobileController(self.ctx) + post_screenshot = await controller.take_screenshot() + except Exception as e: + logger.warning(f"Failed to capture post-action fallback screenshot: {e}") + + if not pre_screenshot or pre_screenshot == "None": + pre_screenshot = post_screenshot + + if not pre_screenshot: + raise ValueError("Failure analysis requires valid screenshots.") + + strategy = self._select_strategy(error_category) + prompt_template_name = strategy.get_prompt_template_name() + tools_declaration = strategy.get_tools() + + try: + llm = get_llm(self.ctx, name="validator_failure_analyzer") + except Exception: + llm = get_llm(self.ctx, name="validator") + + controller = UnifiedMobileController(self.ctx) + executor = MobileActionExecutor(self.ctx, controller) + + # Pre-Action and Post-Action XML formatting + pre_xml_list_str = "Not available." + if hasattr(state, "latest_ui_hierarchy") and state.latest_ui_hierarchy: + try: + width = getattr(self.ctx.device, "device_width", 1080) + height = getattr(self.ctx.device, "device_height", 2400) + pre_xml_list_str, _, _ = format_minimal_list_with_elements( + state.latest_ui_hierarchy, width, height + ) + except Exception as e: + logger.warning(f"Failed to format pre-action XML list: {e}") + + post_xml_list_str = "Not available." + try: + xml_hierarchy = await controller.get_ui_elements() + width = getattr(self.ctx.device, "device_width", 1080) + height = getattr(self.ctx.device, "device_height", 2400) + ocr_results = [] + + fused_xml = fuse_ocr_with_xml(xml_hierarchy, ocr_results) + + post_xml_list_str, elements, _ = format_minimal_list_with_elements( + fused_xml, width, height + ) + state.indexed_points = [el["center"] for el in elements] + state.indexed_elements = elements + except Exception as e: + logger.warning(f"Failed to fetch/format post-action XML list: {e}") + + prompt_path = Path(__file__).parent.joinpath(prompt_template_name) + prompt_template = prompt_path.read_text(encoding="utf-8") + + steps = [] + task_plan = "No task plan yet." + if self.ctx.data_engine: + steps = self.ctx.data_engine.get_agent_friendly_steps() + task_plan_path = get_note_file_path(self.ctx.data_engine.base_dir, "task_plan") + if task_plan_path.exists(): + try: + task_plan = task_plan_path.read_text(encoding="utf-8") + except Exception as e: + task_plan = f"Error reading task plan: {e}" + + plan_and_history = self._build_plan_and_history(steps, task_plan) + failed_step_number = steps[-1]["step_number"] if steps else 1 + failed_action_description = ( + format_action_clean(failed_action) if failed_action else "Planned Action" + ) + + system_prompt = Template(prompt_template).render( + initial_goal=getattr(state, "initial_goal", "Unknown Goal"), + plan_and_history=plan_and_history, + failed_step_number=failed_step_number, + failed_action_description=failed_action_description, + ) + + suffix = strategy.get_user_message_suffix(bool(pre_screenshot), bool(post_screenshot)) + user_msg = ( + suffix.strip() + if suffix + else "Please analyze and repair the failure shown in the last step of the history." + ) + is_pixel_strategy = isinstance(strategy, PixelTargetDisappearedStrategy) + + user_content: list[dict] = [{"type": "text", "text": user_msg}] + if pre_screenshot: + user_content.append( + {"type": "text", "text": "--- Screenshot Seen During System Decision ---"} + ) + user_content.append( + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{pre_screenshot}"}, + } + ) + if not is_pixel_strategy: + user_content.append( + { + "type": "text", + "text": ( + "--- UI Element List Seen During System Decision ---\n" + f"{pre_xml_list_str}" + ), + } + ) + if post_screenshot: + user_content.append( + {"type": "text", "text": "--- Latest Screenshot (Failed State) ---"} + ) + user_content.append( + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{post_screenshot}"}, + } + ) + if not is_pixel_strategy: + user_content.append( + { + "type": "text", + "text": ( + f"--- Latest UI Element List (Failed State) ---\n{post_xml_list_str}" + ), + } + ) + if suffix: + user_content.append( + {"type": "text", "text": f"\n\n--- INSTRUCTIONS ---\n{suffix.strip()}"} + ) + + messages: list[BaseMessage] = [ + SystemMessage(content=system_prompt), + HumanMessage(content=user_content), + ] + + max_iterations = getattr(self, "max_iterations", 15) + iterations = 0 + final_report = None + intermediate_thoughts = [] + + while iterations < max_iterations: + iterations += 1 + self._prune_intermediate_screenshots(messages) + logger.info(f"Iteration {iterations}: Invoking Universal LLM for FailureAnalyzer...") + + if iterations == max_iterations: + messages.append( + HumanMessage( + content=( + "[WARNING] This is your final iteration; only" + " 'report_failure_analysis' is available." + ) + ) + ) + current_tools = [ + t for t in tools_declaration if t.name == "report_failure_analysis" + ] + else: + current_tools = tools_declaration + + bound_llm = llm.bind_tools(current_tools) + + max_retries = 3 + response = None + for attempt in range(max_retries): + try: + response = await bound_llm.ainvoke(messages) + break + except Exception as e: + logger.warning(f"Failure analyzer turn failed on attempt {attempt + 1}: {e}") + if attempt == max_retries - 1: + raise e + await asyncio.sleep(1.0 * (2**attempt)) + + if response is None: + break + + messages.append(response) + + raw_text = response.content if isinstance(response.content, str) else "" + if isinstance(response.content, list): + raw_text = "".join( + b.get("text", "") + for b in response.content + if isinstance(b, dict) and "text" in b + ) + if raw_text: + intermediate_thoughts.append(raw_text) + + tool_calls = response.tool_calls or [] + if not tool_calls: + messages.append( + HumanMessage( + content=( + "CRITICAL: You did not call the `report_failure_analysis` tool." + " You MUST call an action tool to repair or `report_failure_analysis`" + " to finish." + ) + ) + ) + continue + + submit_call = next( + (tc for tc in tool_calls if tc["name"].split(":")[-1] == "report_failure_analysis"), + None, + ) + + if submit_call: + args = submit_call.get("args") or {} + actions = args.get("new_remaining_actions", []) + final_report = { + "status": args.get("status", "cannot_fix"), + "analysis": args.get("analysis", ""), + "new_remaining_actions": actions, + "thoughts": intermediate_thoughts, + } + messages.append( + ToolMessage( + tool_call_id=submit_call.get("id") or str(uuid.uuid4()), + name=submit_call["name"], + content=json.dumps({"status": "acknowledged"}), + status="success", + ) + ) + break + + # Execute other tool calls via MobileActionExecutor + for tc in tool_calls: + name = tc["name"].split(":")[-1] if ":" in tc["name"] else tc["name"] + args = tc.get("args") or {} + tc_id = tc.get("id") or str(uuid.uuid4()) + logger.info(f"Failure analyzer executing tool '{name}'...") + + try: + exec_result = await executor.execute(name, args, tc_id, state) + messages.append(exec_result.to_langchain_tool_message()) + except Exception as e: + logger.error(f"Error executing tool {name}: {e}") + messages.append( + ToolMessage( + tool_call_id=tc_id, + name=name, + content=f"Error executing tool {name}: {e}", + status="error", + ) + ) + + return final_report or { + "status": "cannot_fix", + "analysis": "Max iterations reached without resolution.", + } diff --git a/artemis/agents/validator/failure_analyzer_agent.py b/artemis/agents/validator/failure_analyzer_agent.py new file mode 100644 index 0000000..45ddf91 --- /dev/null +++ b/artemis/agents/validator/failure_analyzer_agent.py @@ -0,0 +1,69 @@ +# 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. + +"""Modular Failure Analyzer Agent adhering to BaseAgent lifecycle.""" + +from typing import Any +from artemis.agents.base import AgentConfig, AgentRegistry, BaseAgent +from artemis.agents.validator.recovery_strategies import RecoveryStrategyRegistry +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +@AgentRegistry.register("failure_analyzer") +class FailureAnalyzerAgent(BaseAgent): + """Specialized agent responsible for diagnosing step failures and triggering autonomous recovery.""" + + def __init__( + self, + config: AgentConfig | None = None, + ctx: Any | None = None, + **kwargs: Any, + ): + super().__init__( + name="failure_analyzer", + role="Autonomous Self-Healing Failure Analyzer", + description="Diagnoses execution exceptions, analyzes screen discrepancies, and runs self-healing sequences.", + config=config, + ctx=ctx, + ) + + async def process(self, state: Any) -> dict[str, Any]: + """Analyzes failure context and executes appropriate recovery strategy.""" + if not self.driver: + raise RuntimeError("FailureAnalyzerAgent requires an active device driver in context.") + + failure_context = getattr(state, "failure_context", {}) or {} + if isinstance(state, dict): + failure_context = state.get("failure_context", {}) + + logger.info(f"FailureAnalyzer diagnosing failure context: {failure_context}") + + recovery_result = await RecoveryStrategyRegistry.attempt_recovery( + driver=self.driver, + failure_context=failure_context, + ) + + if recovery_result and recovery_result.success: + return { + "status": "recovered", + "recovery_action": recovery_result.action_type.value, + "reason": recovery_result.reason, + } + + return { + "status": "unrecoverable", + "reason": "No applicable recovery strategy succeeded.", + } diff --git a/artemis/agents/validator/pixel_safety_net.md b/artemis/agents/validator/pixel_safety_net.md new file mode 100644 index 0000000..d796072 --- /dev/null +++ b/artemis/agents/validator/pixel_safety_net.md @@ -0,0 +1,11 @@ +Role +UI State Verification Expert. Compare Image 1 (Reference) and Image 2 (Current State) to check if the target clicked by the Operator at the red dot is still present and interactive. + +Rules +1. Check Intent: Read the `[Planned Action & Original Thinking]` to identify what Action the Agent intended to perform and why. +2. Specific UI Elements (Priority): If targeting a specific UI control (e.g., button, icon, text, slider handle) โ€” even if it overlays a dynamic video, map, or live stream โ€” focus ONLY on that element's shape. Output `{"is_present": false, ...}` if that exact element has disappeared, shifted away, or is blocked. Ignore background media changes underneath. +3. Blank Space & Backgrounds: If targeting open/blank space (e.g., to wake up hidden controls or dismiss dialogs), ignore natural video/background frame changes. Output `{"is_present": true, "confidence": 1.0}` unless unexpected blocking popups appear. + +Output +Return EXACTLY this raw JSON and nothing else. Keep reasoning strictly under 25 words. +{"reasoning": "brief explanation", "is_present": boolean, "confidence": float} diff --git a/artemis/agents/validator/pixel_target_disappeared_analyzer.md b/artemis/agents/validator/pixel_target_disappeared_analyzer.md new file mode 100644 index 0000000..f446777 --- /dev/null +++ b/artemis/agents/validator/pixel_target_disappeared_analyzer.md @@ -0,0 +1,65 @@ +# Identity & Goal +You are a specialized error recovery agent. You will be provided with a failed action. Your ONLY task is to repair this specific failure and restore the system state so the main system can resume. + + +# Latency Awareness Principle +You operate in an environment with unavoidable "Turn Latency" (a delay of several seconds between screen capture, decision making, and tool execution). +Consequently, any transient UI element (e.g., auto-dismissing overlays, temporary message bars, menus with short timeouts) WILL disappear before your next decision turn if you attempt to interact with them step-by-step. + +# Context +- **Initial Goal**: {{ initial_goal }} +- **History & Plan**: +{{ plan_and_history }} + +# Core Workflow +1. **Analyze**: Compare the two states provided to you: + - **Decision State (`Screenshot Seen During System Decision`)**: This is the screen the system saw when it originally planned the action. It shows the target element successfully located. + - **Failed State (`Latest Screenshot (Failed State)`)**: This is the live screen captured at the moment you were summoned. It shows the current state where the target is missing or unresponsive. + Your goal in this comparison is to: + - Read the Operator's thought process in the decision loop of the failed step to understand what they were originally trying to accomplish with this failed action. + - Understand the expected target layout. + - Identify the specific visual pattern (e.g., a play triangle, a close 'X', an arrow icon) at the failed target area from the reference state. + - Identify why this target element visually disappeared or is no longer responsive (use the **Diagnosis Tips** below). +2. **Recover & Execute**: Restore the target state and execute the failed action. + - **Flexibility**: Focus on outcome over form. Take alternative/adaptive steps (e.g., scroll, navigate, different click paths) if original target positions are invalid. + - **Transient State Handling**: If the failed target is part of a transient state, you MUST NOT execute the recovery and target interaction in separate turns. + - **Atomic Chained Execution**: Combine the summon action (on the Trigger Element) and the target action (on the Transient Target) into a single atomic sequence. Use `click_sequence(sequence=[Trigger, Target])` to execute them with a 50ms delay, bypassing the Turn Latency gap. +3. **Report**: Call `report_failure_analysis` to conclude: + - `status="fixed"`: Successfully executed the failed action or reached the destination state. + - `status="cannot_fix"`: Target is unrecoverable, the app is broken, or the 15-step budget is exceeded. + - `analysis`: Summarize visual mismatches and actions taken. + +# Special Tools +- **`click_sequence(sequence)`**: Designed to defeat Turn Latency. Chain your click sequence to handle transient elements. +- **`ask_explorer(query, context_feedback)`**: Call this visual parsing subagent to search for the visual icon/pattern on the current screen (e.g., `query="play icon"`) if you suspect the target has shifted, or coordinates are ambiguous. + +# Diagnosis Tips +* **Time-Sensitive States (MOST COMMON & PRIMARY CAUSE)**: + - *Cue*: The target element was visible when the system made its decision, but has disappeared in the failed state because it belongs to an auto-closing, short-timeout, or transient UI state. + - *Action*: Re-trigger the state dynamically from the current screen and chain it with the failed step in a rapid consecutive sequence. + - **Mental Model**: Tap the persistent summoner element (Trigger) that wakes up the transient UI + Tap the target coordinates (Failed Target) before it has time to auto-close, combining both into a single atomic movement. + - **Procedure**: + 1. **Locate the Trigger (Target 1)**: Inspect the current screen to identify a persistent parent element or background area (e.g., the container element or parent dropdown header) that, when tapped, forces the transient UI state to reappear. + 2. **Retrieve the Failed Target (Target 2)**: Look up the coordinates of the FAILED action from the failed step in the history (since the target button is currently invisible on your screen). + 3. **Chained Execution**: Call `click_sequence(sequence=[Target 1, Target 2])` to execute both taps consecutively with a 50ms delay, catching the transient element before the timeout closes it. *(Note: This sequence fulfills your workflow requirement to "execute the failed action".)* +* **False Alarm**: + - *Cue*: Target element is actually present on the current screen (minor validation flake). + - *Action*: Directly report it as `"fixed"` and hand control back to the system. +* **Visual Target Shifts**: + - *Cue*: The target visual icon has shifted due to screen layout changes, orientations, or dynamic list updates. + - *Action*: Visually locate the icon on the live screenshot and click its new coordinates. You may call `ask_explorer(query="describe the visual icon shape or labels")` to dynamically locate it and obtain the new coordinates if they are hard to estimate manually. +* **Overlays**: + - *Cue*: Unexpected permission prompts, dialogs, or system popups block the screen. + - *Action*: Tap outside or click the dismiss/close buttons to clear the overlay. +* **Keyboard & Scroll Positioning**: + - *Cue*: Target is scrolled off-screen or blocked by the on-screen keyboard. + - *Action*: Scroll to reposition the target, or hide the keyboard. +* **Navigation Drift**: + - *Cue*: App unexpectedly transitioned to a different tab or screen. + - *Action*: Navigate back to the correct target screen. +* **App Instability**: + - *Cue*: Application crashed, froze, or became unresponsive. + - *Action*: Force stop and relaunch the app. + +# Constraints +The most recent step (Step {{ failed_step_number }}: {{ failed_action_description }}) failed to execute. You must resolve the issue and report the results in 15 interactive steps or fewer. Now please begin the repair analysis. Please note that during the **Recover & Execute** phase, you only need to resolve the currently failed step; do not attempt to execute any subsequent steps or complete the overall goal yourself. Once the failure is resolved, please immediately enter the **Report** phase and return control to the system. diff --git a/artemis/agents/validator/recovery_strategies.py b/artemis/agents/validator/recovery_strategies.py new file mode 100644 index 0000000..446c41a --- /dev/null +++ b/artemis/agents/validator/recovery_strategies.py @@ -0,0 +1,177 @@ +# 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. + +"""Modular Recovery Strategies for the Validation and Self-Healing Engine.""" + +from abc import ABC, abstractmethod +from enum import Enum +from typing import Any +from artemis.drivers.base import BaseDeviceDriver +from artemis.utils.logger import get_logger +from pydantic import BaseModel, Field + +logger = get_logger(__name__) + + +class RecoveryActionType(str, Enum): + RETRY_DYNAMIC_LOCATOR = "retry_dynamic_locator" + DISMISS_DIALOG = "dismiss_dialog" + NAVIGATE_BACK = "navigate_back" + SCROLL_AND_RETRY = "scroll_and_retry" + RESTART_APP = "restart_app" + FALLBACK_COORDINATES = "fallback_coordinates" + + +class RecoveryResult(BaseModel): + """Result of executing an autonomous self-healing recovery strategy.""" + + success: bool + action_type: RecoveryActionType + reason: str + metadata: dict[str, Any] = Field(default_factory=dict) + + +class BaseRecoveryStrategy(ABC): + """Abstract interface for modular self-healing failure recovery tactics.""" + + @property + @abstractmethod + def name(self) -> str: + """Strategy identifier.""" + ... + + @abstractmethod + async def can_handle(self, failure_context: dict[str, Any]) -> bool: + """Determines if this strategy applies to the observed failure mode.""" + ... + + @abstractmethod + async def execute( + self, + driver: BaseDeviceDriver, + failure_context: dict[str, Any], + ) -> RecoveryResult: + """Executes the self-healing corrective actions.""" + ... + + +class DismissDialogRecoveryStrategy(BaseRecoveryStrategy): + """Strategy to detect and dismiss unexpected modal dialogs, popups, or permission prompts.""" + + @property + def name(self) -> str: + return "dismiss_dialog" + + async def can_handle(self, failure_context: dict[str, Any]) -> bool: + category = failure_context.get("category", "") + error_msg = failure_context.get("error_message", "").lower() + return ( + category == "ui_obstruction" + or "permission" in error_msg + or "dialog" in error_msg + or "popup" in error_msg + ) + + async def execute( + self, + driver: BaseDeviceDriver, + failure_context: dict[str, Any], + ) -> RecoveryResult: + logger.info( + "[Recovery] Attempting to dismiss unexpected dialog or overlay via Back button..." + ) + success = await driver.press_key("back") + return RecoveryResult( + success=success, + action_type=RecoveryActionType.DISMISS_DIALOG, + reason="Sent BACK keyevent to dismiss modal overlay", + ) + + +class NavigateBackRecoveryStrategy(BaseRecoveryStrategy): + """Strategy to navigate back when agent navigates to a wrong screen.""" + + @property + def name(self) -> str: + return "navigate_back" + + async def can_handle(self, failure_context: dict[str, Any]) -> bool: + category = failure_context.get("category", "") + return category in ("wrong_screen", "dead_end", "navigation_error") + + async def execute( + self, + driver: BaseDeviceDriver, + failure_context: dict[str, Any], + ) -> RecoveryResult: + logger.info("[Recovery] Wrong screen detected, rolling back navigation...") + success = await driver.press_key("back") + return RecoveryResult( + success=success, + action_type=RecoveryActionType.NAVIGATE_BACK, + reason="Navigated back from incorrect screen", + ) + + +class ScrollAndSearchRecoveryStrategy(BaseRecoveryStrategy): + """Strategy to scroll viewport when target element is off-screen.""" + + @property + def name(self) -> str: + return "scroll_and_search" + + async def can_handle(self, failure_context: dict[str, Any]) -> bool: + category = failure_context.get("category", "") + return category in ("element_not_found", "off_screen", "target_missing") + + async def execute( + self, + driver: BaseDeviceDriver, + failure_context: dict[str, Any], + ) -> RecoveryResult: + direction = failure_context.get("scroll_direction", "down") + logger.info(f"[Recovery] Target missing, scrolling {direction} to reveal element...") + success = await driver.swipe_direction(direction=direction) + return RecoveryResult( + success=success, + action_type=RecoveryActionType.SCROLL_AND_RETRY, + reason=f"Scrolled {direction} to discover offscreen element", + ) + + +class RecoveryStrategyRegistry: + """Central registry and chain-of-responsibility for self-healing recovery strategies.""" + + _strategies: list[BaseRecoveryStrategy] = [ + DismissDialogRecoveryStrategy(), + NavigateBackRecoveryStrategy(), + ScrollAndSearchRecoveryStrategy(), + ] + + @classmethod + def register(cls, strategy: BaseRecoveryStrategy) -> None: + cls._strategies.append(strategy) + + @classmethod + async def attempt_recovery( + cls, + driver: BaseDeviceDriver, + failure_context: dict[str, Any], + ) -> RecoveryResult | None: + """Finds the first applicable strategy and executes recovery.""" + for strategy in cls._strategies: + if await strategy.can_handle(failure_context): + logger.info(f"Applying recovery strategy: {strategy.name}") + return await strategy.execute(driver=driver, failure_context=failure_context) + return None diff --git a/artemis/agents/validator/safety_net.py b/artemis/agents/validator/safety_net.py new file mode 100644 index 0000000..ff33cc8 --- /dev/null +++ b/artemis/agents/validator/safety_net.py @@ -0,0 +1,48 @@ +# 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. + +"""Safety net and system dialog auto-recovery.""" + +from typing import Any +from artemis.drivers.base import BaseDeviceDriver +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +class SafetyNet: + """Detects and resolves unwanted popups, system ANRs, and permission prompts.""" + + SYSTEM_DIALOG_PATTERNS = [ + "is not responding", + "Wait", + "Close app", + "Allow only while using the app", + "While using the app", + "Don't allow", + ] + + @classmethod + async def check_and_recover(cls, screen_data: Any, driver: BaseDeviceDriver) -> bool: + """Inspects screen elements for crash/permission dialogs and dismisses them.""" + for element in screen_data.ui_elements: + text = element.get("text", "") + if any(pattern.lower() in text.lower() for pattern in cls.SYSTEM_DIALOG_PATTERNS): + logger.info( + f"SafetyNet detected system dialog '{text}', attempting auto-dismiss..." + ) + center = element.get("center", [500, 500]) + await driver.tap(center[0], center[1]) + return True + return False diff --git a/artemis/agents/validator/target_disappeared_analyzer.md b/artemis/agents/validator/target_disappeared_analyzer.md new file mode 100644 index 0000000..8bcc61b --- /dev/null +++ b/artemis/agents/validator/target_disappeared_analyzer.md @@ -0,0 +1,61 @@ +# Identity & Goal +You are a specialized error recovery agent. You will be provided with a failed action. Your ONLY task is to repair this specific failure and restore the system state so the main system can resume. + + +# Latency Awareness Principle +You operate in an environment with unavoidable "Turn Latency" (a delay of several seconds between screen capture, decision making, and tool execution). +Consequently, any transient UI element (e.g., auto-dismissing overlays, temporary message bars, menus with short timeouts) WILL disappear before your next decision turn if you attempt to interact with them step-by-step. + +# Context +- **Initial Goal**: {{ initial_goal }} +- **History & Plan**: +{{ plan_and_history }} + +# Core Workflow +1. **Analyze**: Compare the two states provided to you: + - **Decision State (`Screenshot Seen During System Decision`)**: This is the screen the system saw when it originally planned the action. It shows the target element successfully located. + - **Failed State (`Latest Screenshot (Failed State)`)**: This is the live screen captured at the moment you were summoned. It shows the current state where the target is missing or unresponsive. + Your goal in this comparison is to: + - Read the Operator's thought process in the decision loop of the failed step to understand what they were originally trying to accomplish with this failed action. + - Understand the expected target layout. + - Identify why the target element disappeared (use the **Diagnosis Tips** below). +2. **Recover & Execute**: Restore the target state and execute the failed action. + - **Flexibility**: Focus on outcome over form. Take alternative/adaptive steps (e.g., scroll, navigate, different click paths) if original coordinates are invalid. + - **Transient State Handling**: If the failed target is part of a transient state, you MUST NOT execute the recovery and target interaction in separate turns. + - **Atomic Chained Execution**: Combine the summon action (on the Trigger Element) and the target action (on the Transient Target) into a single atomic sequence. Use `click_sequence(sequence=[Trigger, Target])` to execute them with a 50ms delay, bypassing the Turn Latency gap. +3. **Report**: Call `report_failure_analysis` to conclude: + - `status="fixed"`: Successfully executed the failed action or reached the destination state. + - `status="cannot_fix"`: Target is unrecoverable, the app is broken, or the 15-step budget is exceeded. + - `analysis`: Summarize visual mismatches and actions taken. + +# Special Tools +- **`click_sequence(sequence)`**: Designed to defeat Turn Latency. Chain your click sequence to handle transient elements. +- **`ask_explorer(query, context_feedback)`**: Call this visual parsing subagent to search for the visual icon/pattern on the current screen if the target is absent from the standard "UI Element List" but you see it is visually present on the screenshot. + +# Diagnosis Tips +* **Time-Sensitive States (MOST COMMON & PRIMARY CAUSE)**: + - *Cue*: The target element was visible when the system made its decision, but has disappeared in the failed state because it belongs to an auto-closing, short-timeout, or transient UI state. + - *Action*: Re-trigger the state dynamically from the current screen and chain it with the failed step in a rapid consecutive sequence. + - **Mental Model**: Tap the persistent summoner element (Trigger) that wakes up the transient UI + Tap the target coordinates (Failed Target) before it has time to auto-close, combining both into a single atomic movement. + - **Procedure**: + 1. **Locate the Trigger (Target 1)**: Inspect the current screen to identify a persistent parent element or background area (e.g., the container element or parent dropdown header) that, when tapped, forces the transient UI state to reappear. + 2. **Retrieve the Failed Target (Target 2)**: Look up the coordinates of the FAILED action from the failed step in the history (since the target button is currently invisible on your screen). + 3. **Chained Execution**: Call `click_sequence(sequence=[Target 1, Target 2])` to execute both taps consecutively with a 50ms delay, catching the transient element before the timeout closes it. *(Note: This sequence fulfills your workflow requirement to "execute the failed action".)* +* **False Alarm**: + - *Cue*: Target element is actually present on the current screen (minor validation flake). + - *Action*: Directly report it as `"fixed"` and hand control back to the system. +* **Overlays**: + - *Cue*: Unexpected permission prompts, dialogs, or system popups block the screen. + - *Action*: Tap outside or click the dismiss/close buttons to clear the overlay. +* **Keyboard & Scroll Positioning**: + - *Cue*: Target is scrolled off-screen or blocked by the on-screen keyboard. + - *Action*: Scroll to reposition the target, or hide the keyboard. +* **Navigation Drift**: + - *Cue*: App unexpectedly transitioned to a different tab or screen. + - *Action*: Navigate back to the correct target screen. +* **App Instability**: + - *Cue*: Application crashed, froze, or became unresponsive. + - *Action*: Force stop and relaunch the app. + +# Constraints +The most recent step (Step {{ failed_step_number }}: {{ failed_action_description }}) failed to execute. You must resolve the issue and report the results in 15 interactive steps or fewer. Now please begin the repair analysis. Please note that during the **Recover & Execute** phase, you only need to resolve the currently failed step; do not attempt to execute any subsequent steps or complete the overall goal yourself. Once the failure is resolved, please immediately enter the **Report** phase and return control to the system. diff --git a/artemis/agents/validator/tool_declarations.py b/artemis/agents/validator/tool_declarations.py new file mode 100644 index 0000000..de3036e --- /dev/null +++ b/artemis/agents/validator/tool_declarations.py @@ -0,0 +1,1180 @@ +# 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. + +"""Universal Tool Declarations, Execution Results, and Dispatcher for Artemis.""" + +import ast +import asyncio +import base64 +import inspect +import json +import time +from typing import Any, Literal + +from google.genai import types as genai_types +from langchain_core.messages import ToolMessage +from pydantic import BaseModel, Field + +from artemis.config import get_temp_dir, resolve_explorer_version +from artemis.context import ArtemisContext +from artemis.controllers.unified_controller import UnifiedMobileController +from artemis.core.tool_declaration import ToolDeclaration +from artemis.data_engine.trace import trace +from artemis.tools.explorer_tool import _run_explorer_logic +from artemis.graph.state import State +from artemis.tools.mobile.launch_app import find_package +from artemis.utils.app_launch_utils import launch_app_with_retries +from artemis.utils.coordinates import parse_swipe_parameters +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 +from artemis.utils.notes import ( + LIST_NOTES_DOCSTRING, + READ_NOTE_ARG_KEY_DESC, + READ_NOTE_DOCSTRING, + format_list_notes_failure, + format_list_notes_success, + format_read_note_failure, + list_notes_info, + read_note_content, +) + +logger = get_logger(__name__) + +ACTION_TOOL_NAMES = { + "click", + "click_sequence", + "long_press", + "input_text", + "swipe", + "press_key", + "manage_app", + "wait_for_delay", + "wait_for_text", +} + + +class ToolExecutionResult(BaseModel): + """Standardized result produced by executing a mobile automation or utility tool.""" + + tool_call_id: str + tool_name: str + status: Literal["success", "error", "cancelled"] + text_summary: str + screenshot_bytes: bytes | None = None + screenshot_path: str | None = None + ui_elements_text: str | None = None + raw_result: Any = None + metadata: dict[str, Any] = Field(default_factory=dict) + + def to_langchain_tool_message(self) -> ToolMessage: + """Converts execution result into a standardized multimodal LangChain ToolMessage.""" + content_blocks: list[dict[str, Any]] = [ + {"type": "text", "text": self.text_summary or f"Action '{self.tool_name}' completed."} + ] + if self.ui_elements_text: + content_blocks.append( + { + "type": "text", + "text": f"--- UI Element List ---\n{self.ui_elements_text}", + } + ) + if self.screenshot_bytes: + b64_img = base64.b64encode(self.screenshot_bytes).decode("utf-8") + content_blocks.append( + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{b64_img}"}, + } + ) + return ToolMessage( + tool_call_id=self.tool_call_id, + name=self.tool_name, + content=content_blocks, + status=self.status, + ) + + +def prune_intermediate_screenshots(messages: list[Any]) -> None: + """Prunes heavy binary screenshot blocks from intermediate observation messages. + + Keeps the latest screenshot across the conversation history to preserve context budget. + Supports both LangChain BaseMessage and google.genai.types.Content formats. + """ + if not messages: + return + + # Check for google.genai.types.Content objects + if hasattr(messages[0], "parts"): + user_indices_with_images = [] + for idx in range(1, len(messages)): + content = messages[idx] + if getattr(content, "role", None) == "user" and getattr(content, "parts", None): + for part in content.parts: + if getattr(part, "inline_data", None) is not None: + user_indices_with_images.append(idx) + break + + if len(user_indices_with_images) > 1: + indices_to_prune = user_indices_with_images[:-1] + for idx in indices_to_prune: + content = messages[idx] + new_parts = [] + for part in content.parts: + if getattr(part, "inline_data", None) is not None: + try: + new_parts.append( + genai_types.Part.from_text( + text=( + "[Screenshot of intermediate step omitted for performance]" + ) + ) + ) + except Exception: + pass + else: + new_parts.append(part) + content.parts = new_parts + return + + # LangChain BaseMessage format + last_img_msg_idx = -1 + for idx in range(len(messages) - 1, -1, -1): + msg = messages[idx] + msg_content = getattr(msg, "content", None) + if isinstance(msg_content, list): + has_image = any( + isinstance(block, dict) and block.get("type") in ("image_url", "image") + for block in msg_content + ) + if has_image: + last_img_msg_idx = idx + break + + for idx in range(len(messages)): + if idx == last_img_msg_idx: + continue + msg = messages[idx] + msg_content = getattr(msg, "content", None) + if isinstance(msg_content, list): + msg.content = [ + b + for b in msg_content + if not (isinstance(b, dict) and b.get("type") in ("image_url", "image")) + ] + + +async def capture_screenshot_and_parse_ui( + ctx: ArtemisContext, + state: State, + controller: UnifiedMobileController, + skip_settling: bool = False, +) -> tuple[str | None, bytes | None, str | None]: + """Captures screenshot and parses fused XML tree after optional screen settling delay.""" + try: + if skip_settling: + logger.info("Retrieving screen data directly (skip settling)...") + device_data = await controller.get_screen_data() + latest_screenshot_b64 = device_data.base64 + else: + logger.info("Delaying 0.4s before capturing screen state...") + await asyncio.sleep(0.4) + device_data = await controller.get_screen_data() + latest_screenshot_b64 = device_data.base64 + + latest_screenshot_bytes = base64.b64decode(latest_screenshot_b64) + except Exception as e: + logger.error(f"Failed to capture screen state: {e}") + return None, None, 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 + + ocr_results = [] + + fused_xml = fuse_ocr_with_xml(xml_hierarchy, ocr_results) + + minimal_list_str, elements, _ = format_minimal_list_with_elements(fused_xml, width, height) + + state.indexed_points = [el["center"] for el in elements] + state.indexed_elements = elements + except Exception as e: + logger.error(f"Failed to fetch or fuse UI hierarchy: {e}") + minimal_list_str = "Not available due to hierarchy error." + + screenshot_path = None + if ctx.data_engine: + try: + image_name = ctx.data_engine.get_or_create_image( + latest_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(latest_screenshot_bytes) + screenshot_path = str(temp_file) + except Exception as e: + logger.error(f"Failed to save fallback screenshot: {e}") + + return screenshot_path, latest_screenshot_bytes, minimal_list_str + + +def normalize_coordinate_target(target: Any) -> int | list[int] | Any: + """Normalizes coordinate targets or index references to standard types. + + Supports: + - Integer index: e.g. 1 or "1" -> 1 + - Coordinate list: e.g. [500, 280] or ["500", "280"] -> [500, 280] + - Serialized coordinate string: e.g. "[500, 280]" or "500, 280" -> [500, 280] + """ + if isinstance(target, int): + return target + + if isinstance(target, (list, tuple)): + if len(target) == 2: + try: + return [int(target[0]), int(target[1])] + except (ValueError, TypeError): + pass + elif len(target) == 1: + return normalize_coordinate_target(target[0]) + + if isinstance(target, str): + cleaned = target.strip() + if cleaned.isdigit() or (cleaned.startswith("-") and cleaned[1:].isdigit()): + return int(cleaned) + try: + parsed = json.loads(cleaned) + return normalize_coordinate_target(parsed) + except (json.JSONDecodeError, TypeError): + pass + if "," in cleaned: + parts = [p.strip().strip("[]() ") for p in cleaned.split(",")] + if len(parts) == 2: + try: + return [int(parts[0]), int(parts[1])] + except (ValueError, TypeError): + pass + + return target + + +class MobileActionExecutor: + """Encapsulates device action execution, observation capture, and result reporting.""" + + def __init__(self, ctx: ArtemisContext, controller: UnifiedMobileController | None = None): + self.ctx = ctx + self.controller = controller or UnifiedMobileController(ctx) + + @trace(type="tool", name="click") + async def exec_click( + self, + target: list[int] | Any, + state: State, + times: int = 1, + delay_ms: int = 100, + ) -> tuple[str, bytes | None, str | None, str | None]: + width = getattr(self.ctx.device, "device_width", 1080) + height = getattr(self.ctx.device, "device_height", 2400) + try: + target = normalize_coordinate_target(target) + if not isinstance(target, (list, tuple)) or len(target) != 2: + raise ValueError(f"Invalid target format: {target}") + nx, ny = target + x = int(max(0, min(width - 1, int(nx) * width / 1000))) + y = int(max(0, min(height - 1, int(ny) * height / 1000))) + + result = await self.controller.tap_at(x, y, times=times, delay_ms=delay_ms) + err = getattr(result, "error", None) if result else None + success = err is None + outcome = ( + f"Clicked at [{nx}, {ny}] (normalized) successfully." + if success + else f"Error executing click: {err}" + ) + + img_bytes, shot_path, xml_list = None, None, None + if success: + shot_path, img_bytes, xml_list = await capture_screenshot_and_parse_ui( + self.ctx, state, self.controller + ) + except Exception as e: + outcome = f"Error during click: {e}" + img_bytes, shot_path, xml_list = None, None, None + + return outcome, img_bytes, shot_path, xml_list + + @trace(type="tool", name="click_sequence") + async def exec_click_sequence( + self, + sequence: list[int | list[int] | str] | str, + state: State, + delay_ms: int = 50, + ) -> tuple[str, bytes | None, str | None, str | None]: + width = getattr(self.ctx.device, "device_width", 1080) + height = getattr(self.ctx.device, "device_height", 2400) + try: + if isinstance(sequence, str): + sequence_str = sequence.strip() + try: + sequence = json.loads(sequence_str) + except Exception: + try: + sequence = ast.literal_eval(sequence_str) + except Exception: + pass + + if not isinstance(sequence, (list, tuple)): + return ( + f"Error during click sequence: Invalid sequence format: {sequence}", + None, + None, + None, + ) + + outcomes = [] + for i, raw_target in enumerate(sequence): + target = normalize_coordinate_target(raw_target) + if isinstance(target, int): + idx = target + points = getattr(state, "indexed_points", []) or [] + if 1 <= idx <= len(points): + pt = points[idx - 1] + x, y = int(pt[0]), int(pt[1]) + else: + return ( + ( + "Error during click sequence: Index" + f" {idx} is out of range (available: {len(points)})" + ), + None, + None, + None, + ) + elif isinstance(target, (list, tuple)) and len(target) == 2: + nx, ny = target + x = int(max(0, min(width - 1, int(nx) * width / 1000))) + y = int(max(0, min(height - 1, int(ny) * height / 1000))) + else: + return ( + f"Error during click sequence: Invalid target format: {raw_target}", + None, + None, + None, + ) + + result = await self.controller.tap_at(x, y) + err = getattr(result, "error", None) if result else None + if err: + return ( + f"Error executing click at step {i + 1}: {err}", + None, + None, + None, + ) + outcomes.append(f"Tapped at [{x}, {y}]") + if i < len(sequence) - 1: + await asyncio.sleep(max(0, delay_ms) / 1000.0) + + outcome_str = f"Sequence clicked successfully: {'; '.join(outcomes)}" + shot_path, img_bytes, xml_list = await capture_screenshot_and_parse_ui( + self.ctx, state, self.controller + ) + return outcome_str, img_bytes, shot_path, xml_list + except Exception as e: + return f"Error executing click sequence: {e}", None, None, None + + @trace(type="tool", name="long_press") + async def exec_long_press( + self, + target: list[int] | Any, + state: State, + duration_ms: int = 1000, + duration: int | None = None, + ) -> tuple[str, bytes | None, str | None, str | None]: + width = getattr(self.ctx.device, "device_width", 1080) + height = getattr(self.ctx.device, "device_height", 2400) + if duration is not None: + duration_ms = duration + try: + target = normalize_coordinate_target(target) + if not isinstance(target, (list, tuple)) or len(target) != 2: + raise ValueError(f"Invalid target format: {target}") + nx, ny = target + x = int(max(0, min(width - 1, int(nx) * width / 1000))) + y = int(max(0, min(height - 1, int(ny) * height / 1000))) + + result = await self.controller.tap_at( + x, y, long_press=True, long_press_duration=duration_ms + ) + err = getattr(result, "error", None) if result else None + success = err is None + outcome = ( + f"Long pressed at [{nx}, {ny}] (normalized) for {duration_ms}ms successfully." + if success + else f"Error executing long press: {err}" + ) + + img_bytes, shot_path, xml_list = None, None, None + if success: + shot_path, img_bytes, xml_list = await capture_screenshot_and_parse_ui( + self.ctx, state, self.controller + ) + except Exception as e: + outcome = f"Error during long press: {e}" + img_bytes, shot_path, xml_list = None, None, None + + return outcome, img_bytes, shot_path, xml_list + + @trace(type="tool", name="input_text") + async def exec_input_text( + self, + text: str, + target: list[int] | Any, + state: State, + clear_exist: bool = True, + ) -> tuple[str, bytes | None, str | None, str | None]: + width = getattr(self.ctx.device, "device_width", 1080) + height = getattr(self.ctx.device, "device_height", 2400) + try: + if target: + target = normalize_coordinate_target(target) + if not isinstance(target, (list, tuple)) or len(target) != 2: + raise ValueError(f"Invalid target format: {target}") + nx, ny = target + x = int(max(0, min(width - 1, int(nx) * width / 1000))) + y = int(max(0, min(height - 1, int(ny) * height / 1000))) + + await self.controller.tap_at(x, y) + if clear_exist and hasattr(self.controller, "erase_text"): + await self.controller.erase_text() + if hasattr(self.controller, "type_text"): + await self.controller.type_text(text) + + outcome = f"Executed typing '{text}'." + shot_path, img_bytes, xml_list = await capture_screenshot_and_parse_ui( + self.ctx, state, self.controller + ) + except Exception as e: + outcome = f"Error during input text: {e}" + img_bytes, shot_path, xml_list = None, None, None + + return outcome, img_bytes, shot_path, xml_list + + @trace(type="tool", name="swipe") + async def exec_swipe( + self, + action: str | list[int] | Any = None, + duration: int = 400, + state: State | None = None, + duration_ms: int | None = None, + **kwargs: Any, + ) -> tuple[str, bytes | None, str | None, str | None]: + width = getattr(self.ctx.device, "device_width", 1080) + height = getattr(self.ctx.device, "device_height", 2400) + dur = duration_ms if duration_ms is not None else duration + + swipe_input = dict(kwargs) if kwargs else {} + if isinstance(action, dict): + swipe_input.update(action) + elif action is not None: + swipe_input["action"] = action + if "duration" not in swipe_input: + swipe_input["duration"] = dur + + kind, target, final_duration = parse_swipe_parameters(swipe_input, default_duration=dur) + + try: + if kind == "direction": + dir_name = target + dir_map = { + "up": ( + int(width * 0.6), + int(height * 0.8), + int(width * 0.6), + int(height * 0.2), + ), + "down": ( + int(width * 0.6), + int(height * 0.2), + int(width * 0.6), + int(height * 0.8), + ), + "left": ( + int(width * 0.8), + int(height * 0.5), + int(width * 0.2), + int(height * 0.5), + ), + "right": ( + int(width * 0.2), + int(height * 0.5), + int(width * 0.8), + int(height * 0.5), + ), + } + x1, y1, x2, y2 = dir_map[dir_name] + err = await self.controller.swipe_coords(x1, y1, x2, y2, final_duration) + if err: + return f"Error swiping {dir_name}: {err}", None, None, None + outcome = f"Swipe completed successfully. Swiped {dir_name}." + elif kind == "coords": + x1, y1, x2, y2 = target + px1 = int(max(0, min(width - 1, int(x1) * width / 1000))) + py1 = int(max(0, min(height - 1, int(y1) * height / 1000))) + px2 = int(max(0, min(width - 1, int(x2) * width / 1000))) + py2 = int(max(0, min(height - 1, int(y2) * height / 1000))) + err = await self.controller.swipe_coords(px1, py1, px2, py2, final_duration) + if err: + return f"Error dragging: {err}", None, None, None + outcome = f"Swipe completed successfully. Swiped from [{x1}, {y1}] to [{x2}, {y2}]." + else: + return ( + f"Error during swipe: Invalid direction: {action or kwargs}", + None, + None, + None, + ) + + current_state = state or State() + shot_path, img_bytes, xml_list = await capture_screenshot_and_parse_ui( + self.ctx, current_state, self.controller + ) + except Exception as e: + outcome = f"Error during swipe: {e}" + img_bytes, shot_path, xml_list = None, None, None + + return outcome, img_bytes, shot_path, xml_list + + @trace(type="tool", name="press_key") + async def exec_press_key( + self, + key: str, + state: State, + ) -> tuple[str, bytes | None, str | None, str | None]: + try: + valid_keys = {"home", "back", "enter", "delete", "tab", "search", "menu", "app_switch"} + key_str = str(key).lower() + if key_str not in valid_keys: + return f"Error executing key press '{key}'.", None, None, None + + if key_str == "back" and hasattr(self.controller, "go_back"): + await self.controller.go_back() + elif key_str == "home" and hasattr(self.controller, "go_home"): + await self.controller.go_home() + elif key_str == "enter" and hasattr(self.controller, "press_enter"): + await self.controller.press_enter() + elif hasattr(self.controller, "press_key"): + res = await self.controller.press_key(key) + err = getattr(res, "error", None) if res else None + if err: + return f"Error executing key press '{key}': {err}", None, None, None + + outcome = f"Executed key press '{key}'." + shot_path, img_bytes, xml_list = await capture_screenshot_and_parse_ui( + self.ctx, state, self.controller + ) + except Exception as e: + outcome = f"Error during press_key: {e}" + img_bytes, shot_path, xml_list = None, None, None + + return outcome, img_bytes, shot_path, xml_list + + @trace(type="tool", name="manage_app") + async def exec_manage_app( + self, + action: str, + app_name: str, + state: State, + ) -> tuple[str, bytes | None, str | None, str | None]: + try: + res = find_package(self.ctx, app_name, use_fallback=False) + pkg = await res if inspect.iscoroutine(res) else res + if not pkg: + return f"Error finding package for app: {app_name}", None, None, None + target_pkg = pkg + + if action.lower() == "launch": + res_launch = launch_app_with_retries(self.ctx, target_pkg) + if inspect.iscoroutine(res_launch): + res_launch = await res_launch + if isinstance(res_launch, tuple): + success, error_msg = res_launch + else: + success, error_msg = bool(res_launch), "" + + outcome = ( + f"Launched app '{app_name}' ({target_pkg}) successfully." + if success + else f"Failed to launch app '{app_name}': {error_msg}" + ) + elif action.lower() == "stop": + if hasattr(self.controller, "terminate_app"): + res_term = self.controller.terminate_app(target_pkg) + if inspect.iscoroutine(res_term): + await res_term + outcome = f"Terminated app '{app_name}' successfully." + else: + return f"Invalid manage_app action: {action}", None, None, None + + shot_path, img_bytes, xml_list = await capture_screenshot_and_parse_ui( + self.ctx, state, self.controller + ) + except Exception as e: + outcome = f"Error during manage_app: {e}" + img_bytes, shot_path, xml_list = None, None, None + + return outcome, img_bytes, shot_path, xml_list + + @trace(type="tool", name="wait_for_delay") + async def exec_wait_for_delay( + self, + time_in_ms: int, + state: State, + ) -> tuple[str, bytes | None, str | None, str | None]: + try: + delay_s = max(0.0, float(time_in_ms) / 1000.0) + await asyncio.sleep(delay_s) + outcome = f"Waited for {time_in_ms}ms successfully." + shot_path, img_bytes, xml_list = await capture_screenshot_and_parse_ui( + self.ctx, state, self.controller + ) + except Exception as e: + outcome = f"Error during wait_for_delay: {e}" + img_bytes, shot_path, xml_list = None, None, None + + return outcome, img_bytes, shot_path, xml_list + + @trace(type="tool", name="wait_for_text") + async def exec_wait_for_text( + self, + text: str, + wait_state: str | None, + timeout_ms: int | None, + state: State, + ) -> tuple[str, bytes | None, str | None, str | None]: + try: + timeout_s = (timeout_ms or 5000) / 1000.0 + start = time.time() + found = False + target_state = (wait_state or "appear").lower() + + while time.time() - start < timeout_s: + tree = await self.controller.get_ui_elements() + text_present = text.lower() in (tree or "").lower() + if target_state == "appear" and text_present: + found = True + break + elif target_state == "disappear" and not text_present: + found = True + break + await asyncio.sleep(0.5) + + outcome = ( + f"Successfully waited for text '{text}' to {target_state}." + if found + else f"Timed out waiting for text '{text}' to {target_state}." + ) + shot_path, img_bytes, xml_list = await capture_screenshot_and_parse_ui( + self.ctx, state, self.controller + ) + except Exception as e: + outcome = f"Error during wait_for_text: {e}" + img_bytes, shot_path, xml_list = None, None, None + + return outcome, img_bytes, shot_path, xml_list + + @trace(type="tool", name="read_note") + async def exec_read_note( + self, + key: str, + start_line: int | None = None, + end_line: int | None = None, + ) -> str: + try: + base_dir = self.ctx.data_engine.base_dir if self.ctx.data_engine else None + if not base_dir: + return format_read_note_failure(key, "DataEngine not initialized.") + return read_note_content(base_dir, key, start_line, end_line) + except Exception as e: + return format_read_note_failure(key, str(e)) + + @trace(type="tool", name="list_notes") + async def exec_list_notes(self) -> str: + try: + base_dir = self.ctx.data_engine.base_dir if self.ctx.data_engine else None + if not base_dir: + return format_list_notes_failure("DataEngine not initialized.") + notes = list_notes_info(base_dir) + return format_list_notes_success(notes) + except Exception as e: + return format_list_notes_failure(str(e)) + + @trace(type="tool", name="ask_explorer") + async def exec_ask_explorer( + self, + query: str, + context_feedback: str | None, + state: State, + version: str | None = None, + ) -> str: + try: + active_version = resolve_explorer_version( + self.ctx, + explicit_version=version, + agent_or_profile_name="validator", + ) + result = await _run_explorer_logic( + self.ctx, + state, + query, + context_feedback or "", + version=active_version, + ) + if isinstance(result, list): + texts = [ + item["text"] + for item in result + if isinstance(item, dict) and item.get("type") == "text" + ] + return "\n".join(texts) if texts else str(result) + return str(result) + except Exception as e: + return f"Error executing ask_explorer: {e}" + + async def execute( + self, + name: str, + args: dict[str, Any], + tool_call_id: str, + state: State, + ) -> ToolExecutionResult: + """Executes the specified tool by name and wraps result into ToolExecutionResult.""" + raw_name = name.split(":")[-1] if ":" in name else name + res_text = "" + img_bytes, shot_path, xml_list = None, None, None + + if raw_name == "click": + res_text, img_bytes, shot_path, xml_list = await self.exec_click( + args.get("target") or args.get("coordinates") or [], + state, + times=args.get("times", 1), + delay_ms=args.get("delay_ms", 100), + ) + elif raw_name == "click_sequence": + res_text, img_bytes, shot_path, xml_list = await self.exec_click_sequence( + args.get("sequence") or [], + state, + delay_ms=args.get("delay_ms", 50), + ) + elif raw_name == "long_press": + res_text, img_bytes, shot_path, xml_list = await self.exec_long_press( + args.get("target") or args.get("coordinates") or [], + state, + duration_ms=args.get("duration_ms", 1000), + ) + elif raw_name == "input_text": + res_text, img_bytes, shot_path, xml_list = await self.exec_input_text( + args.get("text", ""), + args.get("target") or args.get("coordinates"), + state, + clear_exist=args.get("clear_exist", True), + ) + elif raw_name == "swipe": + res_text, img_bytes, shot_path, xml_list = await self.exec_swipe( + action=args, duration=args.get("duration", 400), state=state + ) + elif raw_name == "press_key": + res_text, img_bytes, shot_path, xml_list = await self.exec_press_key( + args.get("key", "BACK"), state + ) + elif raw_name == "manage_app": + res_text, img_bytes, shot_path, xml_list = await self.exec_manage_app( + args.get("action", "launch"), args.get("app_name", ""), state + ) + elif raw_name == "wait_for_delay": + res_text, img_bytes, shot_path, xml_list = await self.exec_wait_for_delay( + args.get("time_in_ms", 1000), state + ) + elif raw_name == "wait_for_text": + res_text, img_bytes, shot_path, xml_list = await self.exec_wait_for_text( + args.get("text", ""), + args.get("wait_state"), + args.get("timeout_ms", 5000), + state, + ) + elif raw_name == "read_note": + res_text = await self.exec_read_note( + args.get("key", ""), + args.get("start_line"), + args.get("end_line"), + ) + elif raw_name == "list_notes": + res_text = await self.exec_list_notes() + elif raw_name == "ask_explorer": + res_text = await self.exec_ask_explorer( + args.get("query") or args.get("task_description") or "", + args.get("context_feedback"), + state, + ) + else: + res_text = f"Error: Tool '{name}' not supported." + + if shot_path: + state.latest_screenshot = shot_path + + status: Literal["success", "error"] = ( + "success" if (img_bytes is not None or raw_name not in ACTION_TOOL_NAMES) else "error" + ) + if "Error" in res_text or "Failed" in res_text: + status = "error" + + return ToolExecutionResult( + tool_call_id=tool_call_id, + tool_name=name, + status=status, + text_summary=res_text, + screenshot_bytes=img_bytes, + screenshot_path=shot_path, + ui_elements_text=xml_list, + ) + + +CLICK_TOOL = ToolDeclaration( + name="click", + description=( + "[ACTION] Tap/click on the target coordinate on the screen. The screen" + " after click will be returned automatically. For buttons, checkboxes," + " tabs, icons, items in a list, tap ON the element. For text fields /" + " search bars / input boxes, tap INSIDE the input box to focus it." + ), + parameters={ + "type": "object", + "properties": { + "target": { + "type": "array", + "items": {"type": "integer"}, + "description": "Normalized coordinates [x, y] in 0-1000 scale.", + }, + "times": { + "type": "integer", + "description": "Number of taps to perform (default 1). Set to 2 for double click.", + }, + "delay_ms": { + "type": "integer", + "description": "Delay between taps in milliseconds (default 100).", + }, + }, + "required": ["target"], + }, +) + +CLICK_SEQUENCE_TOOL = ToolDeclaration( + name="click_sequence", + description=( + "[ACTION] Executes a sequence of taps one by one in order on the" + " specified targets (e.g. [[500, 280], [885, 362]]). The screen will be returned ONLY after all clicks" + " in the sequence have completed." + ), + parameters={ + "type": "object", + "properties": { + "sequence": { + "type": "array", + "items": { + "type": "array", + "items": {"type": "integer"}, + "description": ( + "Normalized coordinates [x, y] in 0-1000 scale (e.g., [500, 280])," + " or a single integer element index." + ), + }, + "description": "List of targets to tap in sequence, e.g. [[500, 280], [885, 362]].", + }, + "delay_ms": { + "type": "integer", + "description": "Delay between consecutive taps in milliseconds (default 50ms).", + }, + }, + "required": ["sequence"], + }, +) + +LONG_PRESS_TOOL = ToolDeclaration( + name="long_press", + description=( + "[ACTION] Long press on a target coordinate on the screen. The screen" + " after long pressing will be returned automatically." + ), + parameters={ + "type": "object", + "properties": { + "target": { + "type": "array", + "items": {"type": "integer"}, + "description": "Normalized coordinates [x, y] in 0-1000 scale.", + }, + "duration_ms": { + "type": "integer", + "description": "Duration of press in milliseconds (default 1000ms).", + }, + }, + "required": ["target"], + }, +) + +INPUT_TEXT_TOOL = ToolDeclaration( + name="input_text", + description=( + "[ACTION] Type text into an input field on the screen. The screen after" + " typing will be returned automatically. Automatically taps inside the" + " input box at target [x, y] to focus before typing." + ), + parameters={ + "type": "object", + "properties": { + "text": { + "type": "string", + "description": "Text to type into the focused input field.", + }, + "target": { + "type": "array", + "items": {"type": "integer"}, + "description": "Coordinates [x, y] of the input box in 0-1000 scale.", + }, + "clear_exist": { + "type": "boolean", + "description": ( + "Whether to clear existing text in the input box before typing (default True)." + ), + }, + }, + "required": ["text", "target"], + }, +) + +SWIPE_TOOL = ToolDeclaration( + name="swipe", + description=( + "[ACTION] Scroll or swipe on the screen. The screen after swipe will be" + " returned automatically. Set duration >= 1000 to drag-and-drop." + ), + parameters={ + "type": "object", + "properties": { + "action": { + "description": ( + "Swipe action direction string ('up', 'down', 'left', 'right') OR custom" + " coordinates [start_x, start_y, end_x, end_y] in 0-1000 scale." + ) + }, + "duration": { + "type": "integer", + "description": "Swipe duration in milliseconds (default 400).", + }, + }, + "required": ["action"], + }, +) + +PRESS_KEY_TOOL = ToolDeclaration( + name="press_key", + description=( + "[ACTION] Press a physical or virtual system button. The screen after" + " pressing will be returned automatically." + ), + parameters={ + "type": "object", + "properties": { + "key": { + "type": "string", + "description": ( + "Standard Android system button name (ENTER, BACK, HOME, APP_SWITCH)." + ), + } + }, + "required": ["key"], + }, +) + +READ_NOTE_TOOL = ToolDeclaration( + name="read_note", + description=READ_NOTE_DOCSTRING, + parameters={ + "type": "object", + "properties": { + "key": { + "type": "string", + "description": READ_NOTE_ARG_KEY_DESC, + }, + "start_line": { + "type": "integer", + "description": "Start line to read (1-indexed, inclusive).", + }, + "end_line": { + "type": "integer", + "description": "End line to read (1-indexed, inclusive).", + }, + }, + "required": ["key"], + }, +) + +LIST_NOTES_TOOL = ToolDeclaration( + name="list_notes", + description=LIST_NOTES_DOCSTRING, + parameters={"type": "object", "properties": {}}, +) + +MANAGE_APP_TOOL = ToolDeclaration( + name="manage_app", + description=( + "[ACTION] Launch or force stop a specified application. The screen" + " after launch/stop will be returned automatically." + ), + parameters={ + "type": "object", + "properties": { + "action": { + "type": "string", + "description": "The action type ('launch' or 'stop').", + }, + "app_name": { + "type": "string", + "description": "Display name or package name of the application.", + }, + }, + "required": ["action", "app_name"], + }, +) + +WAIT_FOR_DELAY_TOOL = ToolDeclaration( + name="wait_for_delay", + description=( + "[ACTION] Pause execution and wait for a fixed duration in milliseconds." + " The screen after pause will be returned automatically." + ), + parameters={ + "type": "object", + "properties": { + "time_in_ms": { + "type": "integer", + "description": "The exact duration to pause in milliseconds (e.g., 2000).", + } + }, + "required": ["time_in_ms"], + }, +) + +REPORT_FAILURE_ANALYSIS_TOOL = ToolDeclaration( + name="report_failure_analysis", + description=( + "[REPORT] Use this tool to return your final answer when you are done" + " with the analysis and optional repair. This is the ONLY way to return" + " your final conclusion." + ), + parameters={ + "type": "object", + "properties": { + "status": { + "type": "string", + "description": ( + "Whether the failure was repaired locally ('fixed') or not ('cannot_fix')." + ), + }, + "analysis": { + "type": "string", + "description": ( + "Provide the cause analysis and handling method description in a single" + " paragraph, and append the specific operations you ran." + ), + }, + }, + "required": ["status", "analysis"], + }, +) + +REPORT_TASK_STATUS_TOOL = ToolDeclaration( + name="report_task_status", + description=( + "[REPORT] Use this tool to return your final answer when you are done" + " with the task. This is the ONLY way to return your final conclusion." + ), + parameters={ + "type": "object", + "properties": { + "status": { + "type": "string", + "description": ( + "Whether the task was completed successfully ('completed') or" + " failed/unreachable ('failed')." + ), + }, + "explanation": { + "type": "string", + "description": "Provide the final explanation or summary of the task execution.", + }, + }, + "required": ["status", "explanation"], + }, +) + +ASK_EXPLORER_TOOL = ToolDeclaration( + name="ask_explorer", + description=( + "[EXPLORER] Call the Explorer sub-agent to locate an element or" + " coordinate on the screen when XML hierarchy lacks coordinates or" + " visual target is ambiguous." + ), + parameters={ + "type": "object", + "properties": { + "task_description": { + "type": "string", + "description": "The description of what to find on the screen.", + }, + "roi": { + "type": "array", + "items": {"type": "integer"}, + "description": "Optional bounding box [ymin, xmin, ymax, xmax] in 0-1000 scale.", + }, + }, + "required": ["task_description"], + }, +) + +VALIDATOR_TOOLS_DECLARATION: list[ToolDeclaration] = [ + CLICK_TOOL, + LONG_PRESS_TOOL, + INPUT_TEXT_TOOL, + SWIPE_TOOL, + PRESS_KEY_TOOL, + READ_NOTE_TOOL, + LIST_NOTES_TOOL, + MANAGE_APP_TOOL, + WAIT_FOR_DELAY_TOOL, + REPORT_FAILURE_ANALYSIS_TOOL, +] diff --git a/artemis/agents/validator/validator.py b/artemis/agents/validator/validator.py new file mode 100644 index 0000000..427e8e3 --- /dev/null +++ b/artemis/agents/validator/validator.py @@ -0,0 +1,1615 @@ +# 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 asyncio +import base64 +import difflib +import json +import math +import os +from pathlib import Path +import psutil +import re +import signal +import sys +import time +import traceback +import uuid +from uuid import UUID + +from adbutils import AdbClient +from mcp import ClientSession, StdioServerParameters +from mcp.client.stdio import stdio_client + +from langchain_core.messages import HumanMessage, SystemMessage +from artemis.agents.validator.failure_analyzer import ( + FailureAnalyzer, + ValidationErrorCategory, +) +from artemis.clients.ui_automator_client import UIAutomatorClient +from artemis.constants import ( + VALIDATOR_POLL_INTERVAL, + VALIDATOR_POLL_TIMEOUT, + VALIDATOR_UI_HIERARCHY_TIMEOUT, +) +from artemis.context import ArtemisContext, DeviceContext, DevicePlatform +from artemis.controllers.unified_controller import UnifiedMobileController +from artemis.data_engine.trace import CURRENT_TRACE_ID, trace +from artemis.graph.state import State +from artemis.services.llm import get_llm +from artemis.tools.mobile.launch_app import find_package +from artemis.utils import app_launch_utils, image_diff, visualization +from artemis.utils.decorators import wrap_with_callbacks +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + +# Define actions that are permitted to succeed without a detected UI change +ALLOW_NO_UI_CHANGE = { + "wait_for_delay", + "focus_and_clear_text", + "launch_app", + "stop_app", + "press_key", + "back", + "erase_one_char", + "swipe", +} + + +class ValidatorNode: + """Node responsible for executing actions on the device and verifying the results. + + Despite its name, this node handles: + 1. Parsing structured decisions from the Operator. + 2. Executing these actions via an MCP session (calling tools like tap, + swipe, etc.). + 3. Verifying if the action resulted in a UI change. + 4. Triggering failure analysis and local repair if actions fail. + + In a more strict architecture, execution and validation might be split into + separate nodes. + """ + + def __init__(self, ctx: ArtemisContext): + self.ctx = ctx + + def _kill_mcp_server_instantly(self): + """Instantly terminate MCP ADB server child processes using SIGKILL.""" + + try: + current_process = psutil.Process() + children = current_process.children(recursive=True) + for child in children: + try: + cmdline = child.cmdline() + if any("adb_server.py" in part for part in cmdline): + logger.warning(f"Instantly killing MCP Server child process: {child.pid}") + os.kill(child.pid, signal.SIGKILL) + except psutil.NoSuchProcess: + pass + except Exception as e: + logger.error(f"Failed to instantly kill MCP server child processes: {e}") + + async def _get_mcp_session(self): + if getattr(self.ctx, "mcp_session", None) is not None: + return self.ctx.mcp_session + + root_dir = Path(__file__).parent.parent.parent + venv_python = root_dir / ".venv" / "bin" / "python3" + python_exe = str(venv_python) if venv_python.exists() else sys.executable + + server_params = StdioServerParameters( + command=python_exe, + args=[str(root_dir / "mcp" / "adb_server.py")], + env=os.environ.copy(), + ) + + logger.info("Starting persistent MCP server session...") + self.ctx.mcp_client_ctx = stdio_client(server_params) + try: + read, write = await self.ctx.mcp_client_ctx.__aenter__() + self.ctx.mcp_session = ClientSession(read, write) + await self.ctx.mcp_session.__aenter__() + # Handshake timeout protection + await asyncio.wait_for(self.ctx.mcp_session.initialize(), timeout=15.0) + return self.ctx.mcp_session + except Exception as e: + err_stack = traceback.format_exc() + logger.critical( + "CRITICAL: Failed to initialize MCP server session:" + f" {e}\n{err_stack}. Cleaning up child processes." + ) + self._kill_mcp_server_instantly() + self.ctx.mcp_session = None + self.ctx.mcp_client_ctx = None + + logger.warning( + "Falling back to robust local in-process execution bypassing MCP protocol." + ) + + try: + adb = AdbClient(host="localhost", port=5037) + devices = adb.device_list() + if devices: + device_id = devices[0].serial + ui_client = UIAutomatorClient(device_id=device_id) + try: + ui_data = ui_client.get_screen_data() + w, h = ui_data.width, ui_data.height + except Exception: + w, h = 1080, 2400 + local_ctx = ArtemisContext( + trace_id="local-fallback", + device=DeviceContext( + host_platform="LINUX", + mobile_platform=DevicePlatform.ANDROID, + device_id=device_id, + device_width=w, + device_height=h, + ), + adb_client=adb, + ui_adb_client=ui_client, + ) + self._local_controller = UnifiedMobileController(local_ctx) + except Exception as le: + logger.error(f"Failed to setup local fallback controller: {le}") + + return None + + def _parse_decisions(self, structured_decisions: str) -> tuple[list[dict] | None, str | None]: + if not structured_decisions: + return None, "No structured decisions found, nothing to execute." + try: + return json.loads(structured_decisions), None + except json.JSONDecodeError as e: + return None, f"Failed to parse structured decisions: {e}" + + async def _get_initial_screenshot(self, session, state: State) -> tuple[str, str | None]: + screenshot_path = getattr(state, "latest_screenshot", None) + if not screenshot_path: + logger.error("No screenshot path found in state.latest_screenshot") + raise ValueError("No screenshot path found in state.latest_screenshot") + + 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") + except Exception as e: + logger.error(f"Failed to read screenshot from {screenshot_path}: {e}") + raise e + + screenshot_name = Path(screenshot_path).stem + return screenshot_b64, screenshot_name + + async def _poll_for_ui_change( + self, session, pre_screenshot_b64, action_item + ) -> tuple[bool, str, str | None]: + + start_time = time.time() + timeout = VALIDATOR_POLL_TIMEOUT + ui_changed = False + post_screenshot_b64 = None + error_msg = "" + + while time.time() - start_time < timeout: + if session is not None: + result = await session.call_tool("take_screenshot", {}) + success_img, content_img = self._parse_mcp_result(result) + elif getattr(self, "_local_controller", None): + try: + content_img = await self._local_controller.take_screenshot() + success_img = True + except Exception as e: + success_img, content_img = False, str(e) + else: + success_img, content_img = ( + False, + "No session or local controller", + ) + if success_img: + post_screenshot_b64 = content_img + img_before_bytes = base64.b64decode(pre_screenshot_b64) + img_after_bytes = base64.b64decode(post_screenshot_b64) + + ui_changed = image_diff.check_ui_change( + img_before_bytes, img_after_bytes, action_item + ) + if ui_changed: + logger.info("UI change detected during polling.") + break + else: + logger.error(f"Failed to take screenshot via MCP: {content_img}") + + await asyncio.sleep(VALIDATOR_POLL_INTERVAL) + + if not ui_changed: + error_msg = "Screen did not change" + + return ui_changed, error_msg, post_screenshot_b64 + + async def _execute_validation_loop(self, state: State): + actions, error_msg = self._parse_decisions(state.structured_decisions) + if error_msg: + return await state.asanitize_update( + ctx=self.ctx, + update={}, + agent="validator", + ) + + execution = [] + failed_action = None + failure_reason = None + + actions_to_execute = list(actions) + + session = await self._get_mcp_session() + + try: + last_screenshot_b64, last_screenshot_name = await self._get_initial_screenshot( + session, state + ) + # Preserve the turn-initial screenshot (what the operator saw when making decisions) + decision_screenshot_b64 = last_screenshot_b64 + decision_screenshot_name = last_screenshot_name + except Exception: + return await state.asanitize_update( + ctx=self.ctx, + update={}, + agent="validator", + ) + + step_id = None + if state.current_step_id: + step_id = UUID(state.current_step_id) + + analysis_result = {} + + while actions_to_execute: + action_item = actions_to_execute.pop(0) + action_name = action_item.get("action") + + logger.info(f"Executing action: {action_name}") + + action_trace_id = uuid.uuid4() + parent_id = CURRENT_TRACE_ID.get() + token = CURRENT_TRACE_ID.set(action_trace_id) + start_time = time.time() + + if self.ctx.data_engine and step_id: + self.ctx.data_engine.record_trace( + type="action", + name=action_name, + payload={"action": action_item, "status": "running"}, + status="running", + parent_trace_id=parent_id, + step_id=step_id, + trace_id=action_trace_id, + ) + + try: + pre_screenshot_b64 = last_screenshot_b64 + success = False + error_msg = "" + error_category = ValidationErrorCategory.GENERAL + post_screenshot_b64 = None + post_image_name = None + + action_item = dict( + action_item + ) # Make a copy to avoid mutating original state actions + attempts_log = [] + + # 1. Pre-execution validation (Safety Net) + original_coords = list(action_item.get("coordinates") or []) + + is_explorer_candidate = action_item.get("target_class") == "ExplorerCandidate" + is_ocr_element = "[OCR]" in str(action_item.get("target_class") or "") + has_index_metadata = bool( + ( + action_item.get("target_text") + or action_item.get("target_bounds") + or action_item.get("target_resource_id") + ) + and not is_explorer_candidate + and not is_ocr_element + ) + + if has_index_metadata: + # Index-based validation is fastest and safest when metadata is present + ( + validation_passed, + validation_category, + validation_error, + ) = await self._validate_action_precondition(session, action_item, state=state) + + if not validation_passed: + if validation_category == ValidationErrorCategory.XML_BYPASSED: + logger.info( + "Pre-execution XML-based validation bypassed" + f" ({validation_error}). Falling back to" + " Pixel-based validation..." + ) + ( + validation_passed, + validation_category, + validation_error, + ) = await self._validate_action_precondition_pixel( + session, + action_item, + pre_screenshot_b64, + original_coords, + state=state, + ) + else: + logger.info( + "Pre-execution XML-based validation failed:" + f" {validation_error}. Attempting Pixel-based" + " validation fallback..." + ) + ( + pixel_passed, + pixel_category, + pixel_error, + ) = await self._validate_action_precondition_pixel( + session, + action_item, + pre_screenshot_b64, + original_coords, + state=state, + ) + if pixel_passed and pixel_category == ValidationErrorCategory.NONE: + logger.success( + "Pixel-based validation fallback PASSED!" + " Overriding XML validation failure." + ) + validation_passed = True + validation_category = ValidationErrorCategory.NONE + validation_error = "" + else: + # Fall back to Pixel-based validation for pure coordinate/non-index interactions + ( + validation_passed, + validation_category, + validation_error, + ) = await self._validate_action_precondition_pixel( + session, + action_item, + pre_screenshot_b64, + original_coords, + state=state, + ) + + if not validation_passed: + success = False + error_msg = f"Pre-execution validation failed: {validation_error}" + error_category = validation_category + attempts_log.append(error_msg) + + # Capture the live mismatch screenshot to provide context for FailureAnalyzer + try: + if session is not None: + result = await session.call_tool("take_screenshot", {}) + success_img, post_screenshot_b64 = self._parse_mcp_result(result) + elif getattr(self, "_local_controller", None): + post_screenshot_b64 = await self._local_controller.take_screenshot() + success_img = True + else: + success_img, post_screenshot_b64 = False, None + + if success_img and post_screenshot_b64: + if self.ctx.data_engine: + post_image_name = self.ctx.data_engine.get_or_create_image( + base64.b64decode(post_screenshot_b64) + ) + screenshot_path = str( + self.ctx.data_engine.get_image_path(post_image_name) + ) + state.latest_screenshot = screenshot_path + except Exception as e: + logger.error(f"Failed to capture live screenshot for failure analysis: {e}") + post_screenshot_b64 = None + post_image_name = None + else: + # 2. Local execution attempts + max_local_retries = 1 if action_name == "launch_app" else 2 + for attempt in range(max_local_retries): + try: + exec_success, exec_error = await self._exec_action(session, action_item) + except Exception as e: + exec_success = False + exec_error = str(e) + + # Skip polling: assume execution success implies action success + if exec_success: + success = True + error_msg = "" + attempts_log.append("Success") + break + else: + success = False + error_msg = exec_error + attempts_log.append(error_msg) + + if attempt < max_local_retries - 1: + await asyncio.sleep(0.5) + + if not success and not post_screenshot_b64: + logger.info("Action failed, capturing failure screenshot...") + try: + if session is not None: + result = await session.call_tool("take_screenshot", {}) + success_img, post_screenshot_b64 = self._parse_mcp_result(result) + elif getattr(self, "_local_controller", None): + post_screenshot_b64 = await self._local_controller.take_screenshot() + success_img = True + else: + success_img, post_screenshot_b64 = False, None + + if success_img and post_screenshot_b64: + if self.ctx.data_engine: + post_image_name = self.ctx.data_engine.get_or_create_image( + base64.b64decode(post_screenshot_b64) + ) + except Exception as e: + logger.error(f"Failed to capture failure screenshot: {e}") + + # Enrich and record execution + if len(attempts_log) > 1 or not success: + action_item["attempts"] = list(attempts_log) + execution.append(action_item) + + if post_screenshot_b64: + last_screenshot_b64 = post_screenshot_b64 + if self.ctx.data_engine: + post_image_name = self.ctx.data_engine.get_or_create_image( + base64.b64decode(post_screenshot_b64) + ) + last_screenshot_name = post_image_name + + screenshot_path = str(self.ctx.data_engine.get_image_path(post_image_name)) + state.latest_screenshot = screenshot_path + logger.info( + f"Validator updated state.latest_screenshot to: {screenshot_path}" + ) + + # Update action trace with final results + duration = time.time() - start_time + if self.ctx.data_engine and step_id: + relative_time = self.ctx.data_engine.get_relative_time(time.time()) + self.ctx.data_engine.record_trace( + type="action", + name=action_name, + payload={ + "action": action_item, + "success": success, + "error_msg": error_msg, + "post_screenshot": post_image_name, + "timestamp": time.time(), + "relative_time": relative_time, + }, + status="success" if success else "failed", + duration=duration, + parent_trace_id=parent_id, + step_id=step_id, + trace_id=action_trace_id, + ) + finally: + CURRENT_TRACE_ID.reset(token) + + if not success: + failed_action = action_item + failure_reason = error_msg + + logger.warning(f"Action failed: {failed_action}. Triggering failure analysis.") + + analyzer = FailureAnalyzer(self.ctx) + analysis_result = await analyzer.analyze( + state, + failed_action, + failure_reason, + pre_screenshot=decision_screenshot_b64, + post_screenshot=post_screenshot_b64, + pre_screenshot_name=decision_screenshot_name, + post_screenshot_name=post_image_name, + executed_actions=execution[:-1], + unexecuted_actions=list(actions_to_execute), + error_category=error_category, + ) + action_item["repair"] = analysis_result.get("analysis", "No analysis provided.") + + status = analysis_result.get("status") + if status == "fixed": + logger.success("Failure repaired locally!") + actions_to_execute = [] + + # Reload the latest screenshot updated by FailureAnalyzer during repair + try: + screenshot_path = getattr(state, "latest_screenshot", None) + if screenshot_path: + if Path(screenshot_path).exists(): + with open(screenshot_path, "rb") as f: + last_screenshot_b64 = base64.b64encode(f.read()).decode("utf-8") + if self.ctx.data_engine: + last_screenshot_name = self.ctx.data_engine.get_or_create_image( + base64.b64decode(last_screenshot_b64) + ) + logger.info( + "Validator reloaded post-repair screenshot" + f" from: {screenshot_path}" + ) + except Exception as e: + logger.error(f"Failed to reload post-repair screenshot: {e}") + + continue + else: + break + + # Append any skipped actions to the execution log + for skipped in actions_to_execute: + skipped_copy = dict(skipped) + skipped_copy["attempts"] = ["Skipped"] + execution.append(skipped_copy) + + report = { + "execution": execution, + "status": ( + "success" + if success or (failed_action and analysis_result.get("status") == "fixed") + else "failed" + ), + "repair_status": (analysis_result.get("status") if failed_action else None), + } + + # Determine if a distinct post-action screenshot was captured + distinct_post_image_name = ( + last_screenshot_name + if (last_screenshot_name and last_screenshot_name != decision_screenshot_name) + else None + ) + + if self.ctx.data_engine and step_id: + self.ctx.data_engine.update_step_execution_result( + step_id, report, post_image_name=distinct_post_image_name + ) + + return await state.asanitize_update( + ctx=self.ctx, + update={"last_execution_result": report}, + agent="validator", + ) + + @wrap_with_callbacks( + before=lambda: logger.info("Starting Validator Agent..."), + on_success=lambda _: logger.success("Validator Agent"), + on_failure=lambda _: logger.error("Validator Agent"), + ) + @trace(type="agent", name="validator") + async def __call__(self, state: State): + + step_id = None + if state.current_step_id: + step_id = UUID(state.current_step_id) + + original_step_id = getattr(self.ctx.data_engine, "current_step_id", None) + if self.ctx.data_engine and step_id: + self.ctx.data_engine.current_step_id = step_id + + try: + return await self._execute_validation_loop(state) + except asyncio.CancelledError: + logger.warning( + "Validator task cancelled (Stop requested). Instantly killing MCP server processes." + ) + self._kill_mcp_server_instantly() + raise + except Exception as e: + err_stack = traceback.format_exc() + logger.critical(f"CRITICAL ERROR in Validator execution loop: {e}\n{err_stack}") + if self.ctx.data_engine and step_id: + try: + self.ctx.data_engine.update_step_execution_result( + step_id, + { + "status": "error", + "error_msg": str(e), + "traceback": err_stack, + }, + ) + except Exception: + pass + self._kill_mcp_server_instantly() + raise e + finally: + if self.ctx.data_engine: + self.ctx.data_engine.current_step_id = original_step_id + + async def _exec_action( + self, session: ClientSession | None, action_item: dict + ) -> tuple[bool, str]: + action_name = action_item.get("action") + coordinates = action_item.get("coordinates") + + # Intercept non-wait_for_text actions for local fallback + if session is None and getattr(self, "_local_controller", None) is not None: + ctrl = self._local_controller + logger.info(f"Executing in-process action fallback: {action_name}") + try: + if action_name == "tap": + if coordinates and len(coordinates) == 2: + times = action_item.get("times", 1) + delay_ms = action_item.get("delay_ms", 100) + res = await ctrl.tap_at( + coordinates[0], + coordinates[1], + times=times, + delay_ms=delay_ms, + ) + return (True, "Success") if not res.error else (False, res.error) + return False, "Invalid coords" + elif action_name == "long_press_on": + d = action_item.get("duration", 1000) + if coordinates and len(coordinates) == 2: + res = await ctrl.tap_at( + coordinates[0], + coordinates[1], + long_press=True, + long_press_duration=d, + ) + return (True, "Success") if not res.error else (False, res.error) + return False, "Invalid coords" + elif action_name == "swipe": + d = action_item.get("duration", 400) + if coordinates and len(coordinates) == 4: + err = await ctrl.swipe_coords( + coordinates[0], + coordinates[1], + coordinates[2], + coordinates[3], + duration=d, + ) + return (True, "Success") if not err else (False, str(err)) + return False, "Invalid coords" + elif action_name == "focus_and_input_text": + t = action_item.get("text") + clear_before = action_item.get("clear_before_input", False) + if coordinates and len(coordinates) == 2: + await ctrl.tap_at(coordinates[0], coordinates[1]) + + await asyncio.sleep(0.5) + if clear_before: + await ctrl.erase_text() + else: + await ctrl.press_key("123") # Move cursor to end for clean append + success = await ctrl.type_text(t, clear_existing=False) + return (True, "Success") if success else (False, "Failed typing") + return False, "Invalid coords" + elif action_name == "focus_and_clear_text": + if coordinates and len(coordinates) == 2: + await ctrl.tap_at(coordinates[0], coordinates[1]) + + await asyncio.sleep(0.5) + success = await ctrl.erase_text() + return (True, "Success") if success else (False, "Failed erase") + return False, "Invalid coords" + elif action_name == "erase_one_char": + success = await ctrl.erase_text(nb_chars=1) + return (True, "Success") if success else (False, "Failed erase") + elif action_name == "press_key": + keycode = action_item.get("keycode") + success = await ctrl.press_key(keycode) + return (True, "Success") if success else (False, f"Failed press {keycode}") + elif action_name == "back": + success = await ctrl.go_back() + return (True, "Success") if success else (False, "Failed back") + elif action_name == "launch_app": + app_name = action_item.get("app_name") + package_name = await find_package(self.ctx, app_name, use_fallback=False) + if package_name: + success, err_msg = await app_launch_utils.launch_app_with_retries( + ctrl.ctx, package_name + ) + return (True, "Success") if success else (False, err_msg) + return False, "Package not found" + elif action_name == "stop_app": + app_name = action_item.get("app_name") + package_name = await find_package(self.ctx, app_name, use_fallback=False) + if package_name: + success = await ctrl.terminate_app(package_name) + return (True, "Success") if success else (False, "Failed terminate") + return False, "Package not found" + elif action_name == "open_link": + url = action_item.get("url") + success = await ctrl.open_url(url) + return (True, "Success") if success else (False, "Failed open url") + elif action_name == "wait_for_delay": + time_in_ms = action_item.get("time_in_ms", 0) + + await asyncio.sleep(time_in_ms / 1000.0) + return True, "" + else: + return ( + False, + f"Unsupported action for local fallback: {action_name}", + ) + except Exception as e: + return False, f"In-process fallback failed: {e}" + + # Normal MCP Execution Path + try: + if action_name == "tap": + coordinates = action_item.get("coordinates") + times = action_item.get("times") or action_item.get("click_times") or 1 + delay_ms = action_item.get("delay_ms") or action_item.get("delay") or 100 + result = await session.call_tool( + "tap", + { + "coordinates": coordinates, + "times": times, + "delay_ms": delay_ms, + }, + ) + return self._parse_mcp_result(result) + + elif action_name == "long_press_on": + coordinates = action_item.get("coordinates") + duration = action_item.get("duration", 1000) + result = await session.call_tool( + "long_press_on", + {"coordinates": coordinates, "duration": duration}, + ) + return self._parse_mcp_result(result) + + elif action_name == "swipe": + coordinates = action_item.get("coordinates") + duration = action_item.get("duration", 400) + result = await session.call_tool( + "swipe", {"coordinates": coordinates, "duration": duration} + ) + return self._parse_mcp_result(result) + + elif action_name == "focus_and_input_text": + coordinates = action_item.get("coordinates") + text = action_item.get("text", "") + clear_before_input = action_item.get("clear_before_input", False) + result = await session.call_tool( + "focus_and_input_text", + { + "coordinates": coordinates, + "text": text, + "clear_before_input": clear_before_input, + }, + ) + return self._parse_mcp_result(result) + + elif action_name == "focus_and_clear_text": + coordinates = action_item.get("coordinates") + result = await session.call_tool( + "focus_and_clear_text", {"coordinates": coordinates} + ) + return self._parse_mcp_result(result) + + elif action_name == "erase_one_char": + result = await session.call_tool("erase_one_char", {}) + return self._parse_mcp_result(result) + + elif action_name == "press_key": + keycode = action_item.get("keycode", "") + result = await session.call_tool("press_key", {"keycode": keycode}) + return self._parse_mcp_result(result) + + elif action_name == "back": + result = await session.call_tool("back", {}) + return self._parse_mcp_result(result) + + elif action_name == "launch_app": + app_name = action_item.get("app_name") + package_name = await find_package(self.ctx, app_name, use_fallback=False) + if not package_name: + return False, f"Failed to find package for app: {app_name}" + + success, err_msg = await app_launch_utils.launch_app_with_retries( + self.ctx, package_name + ) + return success, "Success" if success else err_msg + + elif action_name == "stop_app": + app_name = action_item.get("app_name") + package_name = await find_package(self.ctx, app_name, use_fallback=False) + if not package_name: + return False, f"Failed to find package for app: {app_name}" + result = await session.call_tool("stop_app", {"package_name": package_name}) + return self._parse_mcp_result(result) + + elif action_name == "open_link": + url = action_item.get("url") + result = await session.call_tool("open_link", {"url": url}) + return self._parse_mcp_result(result) + + elif action_name == "wait_for_delay": + time_in_ms = action_item.get("time_in_ms", 0) + logger.info(f"Validator performing wait_for_delay for {time_in_ms}ms...") + + await asyncio.sleep(time_in_ms / 1000.0) + return True, "" + + else: + return False, f"Unsupported action: {action_name}" + except Exception as e: + return False, f"MCP call failed: {e}" + + @trace(type="tool", name="safety_net_validation") + async def _validate_action_precondition( + self, session, action_item: dict, state: State | None = None + ) -> tuple[bool, ValidationErrorCategory, str]: + """Safety Net: Validates if the target element for interactive actions + + is still present on the live screen and is substantially consistent, + with a short retry loop to accommodate screen transitions and loading + delays. + """ + max_pre_attempts = 3 + pre_retry_delay = 0.4 + + passed = False + category = ValidationErrorCategory.NONE + reason = "" + + for pre_attempt in range(1, max_pre_attempts + 1): + passed, category, reason = await self._validate_action_precondition_single( + session, action_item, state + ) + if passed: + return True, category, reason + + if category == ValidationErrorCategory.XML_BYPASSED: + # Do not retry on XML bypass/timeout; return immediately + # so it falls back to Pixel-based VLM validation + return False, category, reason + + if pre_attempt < max_pre_attempts: + logger.info( + "Pre-execution validation failed on attempt" + f" {pre_attempt}/{max_pre_attempts}: {reason}. Retrying in" + f" {pre_retry_delay}s..." + ) + await asyncio.sleep(pre_retry_delay) + else: + break + + return passed, category, reason + + async def _validate_action_precondition_single( + self, session, action_item: dict, state: State | None = None + ) -> tuple[bool, ValidationErrorCategory, str]: + """Safety Net: Validates if the target element for interactive actions + + is still present on the live screen and is substantially consistent. + + Returns (True, ValidationErrorCategory.NONE, "") if validation passes or + is skipped. + Returns (False, category, error_msg) if validation fails. + """ + action_name = action_item.get("action") + + # 1. White-list filtering: Only validate interactive, element-focused actions + if action_name not in [ + "tap", + "long_press_on", + "focus_and_input_text", + "focus_and_clear_text", + ]: + return True, ValidationErrorCategory.NONE, "" + + target_text = action_item.get("target_text") + target_bounds = action_item.get("target_bounds") # [l, t, r, b] + target_resource_id = action_item.get("target_resource_id") + + width = 1080 + 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): + width = w_raw + height = h_raw + else: + device = getattr(self.ctx, "device", None) + if device: + width = getattr(device, "device_width", 1080) + height = getattr(device, "device_height", 2400) + + # Calculate diagonal-based scale factor for resolution-independent thresholds + + ref_diagonal = math.sqrt(1080**2 + 2400**2) # ~2631.8 + current_diagonal = math.sqrt(width**2 + height**2) + scale_factor = current_diagonal / ref_diagonal + + # 2. Skip validation if there's no target metadata to match against + if not target_text and not target_bounds and not target_resource_id: + return True, ValidationErrorCategory.NONE, "" + + logger.info( + f"Pre-execution validation: checking target '{target_text}'" + f" ({target_resource_id}) at {target_bounds}..." + ) + + # 3. Pull live XML tree with strict timeout to avoid hanging the execution flow + + try: + elements = None + if session is not None: + try: + result = await asyncio.wait_for( + session.call_tool("get_ui_hierarchy", {}), + timeout=VALIDATOR_UI_HIERARCHY_TIMEOUT, + ) + if hasattr(result, "content") and result.content: + text_content = ( + result.content[0].text + if hasattr(result.content[0], "text") + else str(result.content[0]) + ) + elements = json.loads(text_content) + except Exception as e: + logger.warning( + f"Failed to get live XML via MCP: {e}. Falling back to" + " Pixel-based validation." + ) + return ( + False, + ValidationErrorCategory.XML_BYPASSED, + f"XML hierarchy fetch timed out or errored: {e}", + ) + elif getattr(self, "_local_controller", None): + try: + elements = await asyncio.wait_for( + self._local_controller.get_ui_elements(), + timeout=VALIDATOR_UI_HIERARCHY_TIMEOUT, + ) + except Exception as e: + logger.warning( + f"Failed to get live XML via local controller: {e}." + " Falling back to Pixel-based validation." + ) + return ( + False, + ValidationErrorCategory.XML_BYPASSED, + f"XML hierarchy fetch timed out or errored: {e}", + ) + + if not elements or not isinstance(elements, list): + logger.warning( + "Live hierarchy empty or invalid. Falling back to Pixel-based validation." + ) + return ( + False, + ValidationErrorCategory.XML_BYPASSED, + "Live hierarchy empty or invalid.", + ) + + except Exception as e: + logger.error( + f"Unexpected error during live XML fetch: {e}. Falling back to" + " Pixel-based validation." + ) + return ( + False, + ValidationErrorCategory.XML_BYPASSED, + f"Unexpected error during live XML fetch: {e}", + ) + + # 4. Matching Heuristics & Scoring + + # Normalization Helper + def normalize_text(s: str | None) -> str: + if not s: + return "" + s = s.lower() + s = re.sub(r"\(\d+\+?\)", "", s) + s = re.sub(r"\[\d+\+?\]", "", s) + s = re.sub(r"[^\w\s\d]", "", s) + return s.strip() + + # Recursive text aggregator for candidate container nodes + def aggregate_text(elem: dict) -> str: + text = ( + elem.get("text") or elem.get("content-desc") or elem.get("accessibilityText") or "" + ) + if not isinstance(text, str): + text = str(text) + children = elem.get("children") or [] + if children: + for child in children: + if isinstance(child, dict): + text += " " + aggregate_text(child) + return text + + norm_target_text = normalize_text(target_text) + + candidates = [] + for elem in elements: + if not isinstance(elem, dict): + continue + + elem_bounds_str = elem.get("bounds") + elem_bounds = visualization.parse_bounds(elem_bounds_str) + if not elem_bounds: + continue + + elem_raw_text = aggregate_text(elem) + elem_norm_text = normalize_text(elem_raw_text) + elem_res_id = elem.get("resource-id") or elem.get("resourceId") or "" + + elem_cx, elem_cy = visualization.get_center_coordinates(*elem_bounds) + + original_coords = action_item.get("coordinates") or [0, 0] + dist = math.sqrt( + (elem_cx - original_coords[0]) ** 2 + (elem_cy - original_coords[1]) ** 2 + ) + + id_match = target_resource_id and elem_res_id and target_resource_id == elem_res_id + + text_score = 0.0 + if norm_target_text and elem_norm_text: + if norm_target_text == elem_norm_text: + text_score = 1.0 + elif norm_target_text in elem_norm_text or elem_norm_text in norm_target_text: + text_score = 0.8 + else: + text_score = difflib.SequenceMatcher( + None, norm_target_text, elem_norm_text + ).ratio() + + # Calculate bounds match (IoU) + bounds_score = 0.0 + size_mismatch = False + if target_bounds and elem_bounds and len(target_bounds) == 4 and len(elem_bounds) == 4: + l1, t1, r1, b1 = target_bounds + l2, t2, r2, b2 = elem_bounds + + # Check size ratio mismatch to prevent matching small widget to giant container + w1, h1 = r1 - l1, b1 - t1 + w2, h2 = r2 - l2, b2 - t2 + if w1 > 0 and h1 > 0 and w2 > 0 and h2 > 0: + size_ratio_w = max(w1, w2) / min(w1, w2) + size_ratio_h = max(h1, h2) / min(h1, h2) + if size_ratio_w > 2.5 or size_ratio_h > 2.5: + size_mismatch = True + + li = max(l1, l2) + ti = max(t1, t2) + ri = min(r1, r2) + bi = min(b1, b2) + if ri > li and bi > ti: + inter_area = (ri - li) * (bi - ti) + union_area = (r1 - l1) * (b1 - t1) + (r2 - l2) * (b2 - t2) - inter_area + if union_area > 0: + bounds_score = inter_area / union_area + + # Calculate Coordinate Containment + contains_coord = False + if ( + original_coords + and len(original_coords) == 2 + and elem_bounds + and len(elem_bounds) == 4 + ): + cx, cy = original_coords + l2, t2, r2, b2 = elem_bounds + contains_coord = l2 <= cx <= r2 and t2 <= cy <= b2 + + # Apply size mismatch penalty + if size_mismatch and not id_match: + text_score = 0.0 + bounds_score = 0.0 + contains_coord = False + + # Define signal weights + W_id = 0.5 + W_text = 0.4 + W_bounds = 0.3 + W_coord = 0.3 + + signals = [] + weights = [] + identity_signals = [] + identity_weights = [] + + # 1. Resource ID Signal + if target_resource_id: + weights.append(W_id) + identity_weights.append(W_id) + if id_match: + signals.append(W_id * 1.0) + identity_signals.append(W_id * 1.0) + elif not elem_res_id: + signals.append(W_id * 0.3) # Soft penalty for missing ID + identity_signals.append(W_id * 0.3) + else: + signals.append(W_id * -0.5) # Hard penalty for mismatched ID + identity_signals.append(W_id * -0.5) + + # 2. Text Signal + if target_text: + weights.append(W_text) + identity_weights.append(W_text) + if norm_target_text and elem_norm_text: + signals.append(W_text * text_score) + identity_signals.append(W_text * text_score) + elif not elem_norm_text: + signals.append(W_text * 0.0) # Penalty for missing expected text + identity_signals.append(W_text * 0.0) + + # 3. Bounds Signal + if target_bounds: + weights.append(W_bounds) + signals.append(W_bounds * bounds_score) + + # 4. Coordinate Containment Signal + if original_coords: + weights.append(W_coord) + signals.append(W_coord * (1.0 if contains_coord else 0.0)) + + if weights: + score = sum(signals) / sum(weights) + else: + score = 1.0 + + if identity_weights: + identity_score = sum(identity_signals) / sum(identity_weights) + else: + identity_score = 1.0 + + decay = max(0.5, 1.0 - (dist / 800.0)) + score *= decay + + candidates.append( + { + "element": elem, + "center": [elem_cx, elem_cy], + "bounds": elem_bounds, + "text": elem_raw_text, + "resource_id": elem_res_id, + "score": score, + "identity_score": identity_score, + "distance": dist, + "text_score": text_score, + "id_match": id_match, + } + ) + + if not candidates: + logger.warning( + "No candidates with valid bounds found. Falling back to Pixel-based validation." + ) + return ( + False, + ValidationErrorCategory.XML_BYPASSED, + "No candidates with valid bounds found.", + ) + + candidates.sort(key=lambda x: x["score"], reverse=True) + best = candidates[0] + + passed = False + category = ValidationErrorCategory.NONE + reason = "" + + threshold = 0.55 if best["distance"] <= 150 * scale_factor else 0.75 + + if best["score"] >= threshold: + passed = True + if best["distance"] <= 200 * scale_factor: + new_coords = best["center"] + old_coords = action_item.get("coordinates") + if old_coords != new_coords: + logger.success( + "Pre-execution validation SUCCESS" + f" (score={best['score']:.2f}). Target element" + f" '{target_text}' shifted by {best['distance']:.1f}px." + f" Self-healing: correcting coordinates {old_coords} ->" + f" {new_coords}." + ) + action_item["coordinates"] = new_coords + else: + logger.info( + "Pre-execution validation SUCCESS" + f" (score={best['score']:.2f}). Target element matched but" + f" shifted by {best['distance']:.1f}px" + f" (>{200 * scale_factor:.1f}px). Bypassing coordinate" + " self-healing to remain conservative." + ) + else: + original_coords = action_item.get("coordinates") + + # 1. Shifted Case: we must be conservative to avoid confusing disappearance as shift. + # Shift misclassification as disappeared is acceptable, but disappeared as shift + # causes errors. + is_valid_shift = False + + # Enforce strict shift criteria (preferring disappeared over shifted) + max_allowed_distance = 100.0 * scale_factor + if best.get("id_match", False): + # Allow larger distance threshold only for strong resource ID matches + max_allowed_distance = 300.0 * scale_factor + + if best["distance"] <= max_allowed_distance: + if best["identity_score"] >= 0.85: + # High confidence identity match (e.g. exact text or exact ID) + is_valid_shift = True + elif best["identity_score"] >= 0.5 and best.get("id_match", False): + # Moderate confidence with exact resource ID match + is_valid_shift = True + + if is_valid_shift: + category = ValidationErrorCategory.TARGET_SHIFTED + reason = ( + f"Target element '{target_text}' ({target_resource_id})" + f" expected at {original_coords} (bounds: {target_bounds})" + f" has shifted.\n- New location: {best['center']} (bounds:" + f" {best['bounds']})" + ) + else: + # 2. Find the most specific element containing expected coordinates + occupant = None + if original_coords and len(original_coords) == 2: + orig_cx, orig_cy = original_coords + for elem in elements: + if not isinstance(elem, dict): + continue + elem_bounds_str = elem.get("bounds") + elem_bounds = visualization.parse_bounds(elem_bounds_str) + if elem_bounds: + left, top, right, bottom = elem_bounds + if left <= orig_cx <= right and top <= orig_cy <= bottom: + elem_area = (right - left) * (bottom - top) + if not occupant or elem_area < occupant[2]: + occupant = (elem, elem_bounds, elem_area) + + def is_interactive(e: dict) -> bool: + clickable = e.get("clickable") + long_clickable = e.get("long-clickable") or e.get("longClickable") + focusable = e.get("focusable") + enabled = e.get("enabled") + + def to_bool(v) -> bool: + if v is None: + return False + if isinstance(v, bool): + return v + return str(v).lower() == "true" + + return ( + to_bool(clickable) + or to_bool(long_clickable) + or to_bool(focusable) + or to_bool(enabled) + ) + + is_occupied = False + occupant_desc = "anonymous element" + occupant_bounds = None + + if occupant: + elem, elem_bounds, _ = occupant + occupant_text = ( + elem.get("text") + or elem.get("content-desc") + or elem.get("accessibilityText") + or "" + ) + occupant_res_id = elem.get("resource-id") or elem.get("resourceId") or "" + occupant_bounds = elem_bounds + + # Check if the occupant has content + has_content = bool(occupant_text or occupant_res_id) + + # Check if this occupant or any container covering this point is interactive + coords_interactive = False + if original_coords and len(original_coords) == 2: + orig_cx, orig_cy = original_coords + for e in elements: + if not isinstance(e, dict): + continue + eb_str = e.get("bounds") + eb = visualization.parse_bounds(eb_str) + if eb: + e_left, e_top, e_right, e_bottom = eb + if e_left <= orig_cx <= e_right and e_top <= orig_cy <= e_bottom: + # Exclude full-screen root background elements + e_width = e_right - e_left + e_height = e_bottom - e_top + is_full_screen = ( + e_width >= width - 10 and e_height >= height - 10 + ) + + if is_interactive(e) and not is_full_screen: + coords_interactive = True + break + + # Exclude full-screen or giant parent containers from being considered blockers + is_occupant_full_screen = False + if occupant_bounds: + o_w = occupant_bounds[2] - occupant_bounds[0] + o_h = occupant_bounds[3] - occupant_bounds[1] + is_occupant_full_screen = o_w >= width - 10 and o_h >= height - 10 + + # If target bounds are known, check if occupant is a giant container + if target_bounds: + t_w = target_bounds[2] - target_bounds[0] + t_h = target_bounds[3] - target_bounds[1] + t_area = t_w * t_h + o_area = o_w * o_h + if t_area > 0 and o_area > 3.0 * t_area: + is_occupant_full_screen = True + + if (has_content or coords_interactive) and not is_occupant_full_screen: + is_occupied = True + occupant_desc = ( + f"'{occupant_text}' ({occupant_res_id})" + if has_content + else "interactive anonymous element" + ) + + # If coordinates are occupied by a different active/labeled element + if is_occupied: + category = ValidationErrorCategory.TARGET_OCCUPIED + reason = ( + f"The expected position {original_coords} for" + f" '{target_text}' ({target_resource_id}) is" + " occupied/intercepted by a different element:" + f" {occupant_desc} at bounds {occupant_bounds}." + ) + else: + # 3. Completely Missing Case + category = ValidationErrorCategory.TARGET_DISAPPEARED + reason = ( + f"Target element '{target_text}' ({target_resource_id})" + f" expected at {original_coords} (bounds:" + f" {target_bounds}) was not found on the screen." + ) + + return passed, category, reason + + def _parse_mcp_result(self, result) -> tuple[bool, str]: + """Helper to parse MCP tool result.""" + if hasattr(result, "content") and result.content: + text_content = ( + result.content[0].text + if hasattr(result.content[0], "text") + else str(result.content[0]) + ) + if text_content.startswith("Error:") or text_content.startswith("Failed:"): + return False, text_content + elif text_content == "Success": + return True, "" + elif text_content == "Failed": + return False, "Action failed on server" + return True, text_content + return True, "" + + @trace(type="tool", name="safety_net_pixel_validation") + async def _validate_action_precondition_pixel( + self, + session, + action_item: dict, + pre_screenshot_b64: str, + original_coords: list[int] | None = None, + state: State | None = None, + ) -> tuple[bool, ValidationErrorCategory, str]: + """Pixel-level Safety Net: Validates if visual target button/component at + coordinates is still present, visible, and interactive on the live screen + compared to the original screenshot. + """ + action_name = action_item.get("action") + + # 1. White-list filtering: Only validate interactive, element-focused actions + if action_name not in [ + "tap", + "long_press_on", + "focus_and_input_text", + "input_text", + "click_coordinate", + "click", + "long_press", + ]: + return True, ValidationErrorCategory.PIXEL_BYPASSED, "" + + current_coords = action_item.get("coordinates") + orig_coords = original_coords or current_coords + + if not orig_coords or len(orig_coords) != 2: + logger.info("Pixel safety net skipped: no valid coordinates.") + return True, ValidationErrorCategory.PIXEL_BYPASSED, "" + + if not current_coords or len(current_coords) != 2: + current_coords = orig_coords + + logger.info( + "Pre-execution pixel validation: checking target position" + f" {orig_coords} in original image vs {current_coords} on live" + " screen..." + ) + + # 2. Decode and crop the static reference (original) image ONCE outside the loop to save CPU + + try: + orig_bytes = base64.b64decode(pre_screenshot_b64) + orig_crop_bytes = visualization.crop_and_annotate_target( + orig_bytes, orig_coords, crop_size=None, dot_radius=15 + ) + except Exception as e: + logger.error( + f"Failed to prepare original crop for pixel safety net: {e}. Bypassing check." + ) + return True, ValidationErrorCategory.PIXEL_BYPASSED, "" + + # 3. Get LLM Configuration once outside the loop + + try: + try: + llm = get_llm(self.ctx, name="validator_pixel_safety_net") + except Exception: + llm = get_llm(self.ctx, name="validator") + + # Load prompt template once + prompt_path = Path(__file__).parent.joinpath("pixel_safety_net.md") + prompt = prompt_path.read_text(encoding="utf-8") + except Exception as e: + logger.error( + f"Failed to initialize LLM/prompt for pixel safety net: {e}. Bypassing check." + ) + return True, ValidationErrorCategory.PIXEL_BYPASSED, "" + + # 4. Start high-precision execution validation retry loop + max_attempts = 3 + retry_delay_seconds = 0.3 + last_error = None + + for attempt in range(1, max_attempts + 1): + logger.info(f"Pre-execution pixel validation: Attempt {attempt}/{max_attempts}...") + try: + # A. Take fresh live screenshot (captures state changes / dynamic loading settling) + live_screenshot_b64 = None + if session is not None: + result = await session.call_tool("take_screenshot", {}) + success_img, live_screenshot_b64 = self._parse_mcp_result(result) + elif getattr(self, "_local_controller", None): + live_screenshot_b64 = await self._local_controller.take_screenshot() + success_img = True + else: + success_img, live_screenshot_b64 = False, None + + if not success_img or not live_screenshot_b64: + raise Exception("Failed to acquire live screenshot from controller/MCP tool.") + + # B. Decode and crop the live screen target + live_bytes = base64.b64decode(live_screenshot_b64) + live_crop_bytes = visualization.crop_and_annotate_target( + live_bytes, current_coords, crop_size=None, dot_radius=15 + ) + + orig_b64 = base64.b64encode(orig_crop_bytes).decode("utf-8") + live_b64 = base64.b64encode(live_crop_bytes).decode("utf-8") + + # C. Formulate multi-modal verification message + user_content = [ + {"type": "text", "text": "[Image 1 (Reference)]"}, + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{orig_b64}"}, + }, + {"type": "text", "text": "[Image 2 (Current State)]"}, + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{live_b64}"}, + }, + ] + if state: + thoughts = [] + native_thought = getattr(state, "operator_native_thinking", None) + if native_thought and native_thought.strip(): + thoughts.append(native_thought.strip()) + raw_thought = getattr(state, "operator_raw_thinking", None) + if raw_thought and raw_thought.strip(): + thoughts.append(raw_thought.strip()) + + operator_thought = "\n\n".join(thoughts) + if operator_thought.strip(): + context_text = ( + "[Planned Action & Original Thinking]\nAction:" + f" {action_name}\nOriginal" + f" Thinking:\n{operator_thought.strip()}" + ) + user_content.append({"type": "text", "text": f"\n{context_text.strip()}"}) + + messages = [ + SystemMessage(content=prompt), + HumanMessage(content=user_content), + ] + + # D. Invoke Universal VLM + response = await llm.ainvoke(messages) + output = response.content if isinstance(response.content, str) else "" + if isinstance(response.content, list): + output = "".join( + b.get("text", "") + for b in response.content + if isinstance(b, dict) and "text" in b + ) + + if output.startswith("```"): + code_match = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", output) + if code_match: + output = code_match.group(1).strip() + + res_json = json.loads(output) + reasoning = res_json.get("reasoning", "") + is_present = res_json.get("is_present", True) + confidence = res_json.get("confidence", 1.0) + + logger.info( + f"Pixel validation attempt {attempt}/{max_attempts} result:" + f" is_present={is_present}, confidence={confidence:.2f}, reasoning={reasoning}" + ) + + # E. Branching based on outcome + if is_present: + logger.info(f"Pixel validation SUCCESS on attempt {attempt}/{max_attempts}.") + return True, ValidationErrorCategory.NONE, "" + + # Target not present: fail or bypass based on confidence without retry! + if confidence >= 0.7: + reason_msg = ( + f"Pixel-level validation failed: {reasoning}" + if reasoning + else ( + "Pixel-level validation failed: The target UI" + " element at the exact coordinates has disappeared," + " changed, or is blocked." + ) + ) + return ( + False, + ValidationErrorCategory.PIXEL_TARGET_DISAPPEARED, + reason_msg, + ) + else: + logger.warning( + "Pixel validation failed but confidence" + f" ({confidence:.2f}) is below 0.7. Bypassing check." + ) + return True, ValidationErrorCategory.PIXEL_BYPASSED, "" + + except Exception as attempt_err: + logger.warning( + f"Error during pixel validation attempt {attempt}/{max_attempts}: {attempt_err}" + ) + last_error = attempt_err + if attempt < max_attempts: + await asyncio.sleep(retry_delay_seconds) + continue + else: + logger.error( + "Pixel validation repeatedly failed with errors. Last" + f" error: {last_error}. Bypassing check to avoid" + " blocking flow." + ) + return True, ValidationErrorCategory.PIXEL_BYPASSED, "" + + return True, ValidationErrorCategory.PIXEL_BYPASSED, "" diff --git a/artemis/agents/video_analyzer/__init__.py b/artemis/agents/video_analyzer/__init__.py new file mode 100644 index 0000000..76497fb --- /dev/null +++ b/artemis/agents/video_analyzer/__init__.py @@ -0,0 +1,15 @@ +# 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. + +"""Video analyzer utility for analyzing video content with Gemini models.""" diff --git a/artemis/agents/video_analyzer/conflict_resolution.py b/artemis/agents/video_analyzer/conflict_resolution.py new file mode 100644 index 0000000..8e5140b --- /dev/null +++ b/artemis/agents/video_analyzer/conflict_resolution.py @@ -0,0 +1,137 @@ +# 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 logging + +logger = logging.getLogger(__name__) + + +class ConflictResolutionService: + @staticmethod + def calculate_1d_iou(a, b): + start = max(a[0], b[0]) + end = min(a[1], b[1]) + intersection = max(0.0, end - start) + union = (a[1] - a[0]) + (b[1] - b[0]) - intersection + return intersection / (union + 1e-6) + + @staticmethod + def calculate_1d_iomin(a, b): + start = max(a[0], b[0]) + end = min(a[1], b[1]) + intersection = max(0.0, end - start) + min_len = min(a[1] - a[0], b[1] - b[0]) + if min_len <= 0: + return 0.0 + return intersection / min_len + + @staticmethod + def sanitize_interval(interval): + start = interval.get("start", 0.0) + end = interval.get("end", "unknown") + + if isinstance(start, (int, float)): + start = max(0.0, float(start)) + + if end != "unknown" and isinstance(end, (int, float)): + end = max(0.0, float(end)) + if start > end: + start, end = end, start + + interval["start"] = start + interval["end"] = end + + if "confidence_score" not in interval or interval["confidence_score"] is None: + interval["confidence_score"] = 0.0 + + if end != "unknown": + interval["duration"] = float(end) - float(start) + else: + interval["duration"] = 0.0 + + return interval + + @staticmethod + def group_by_semantic_class(proposals): + groups = {} + for p in proposals: + target = p.get("target", "unknown") + if target not in groups: + groups[target] = [] + groups[target].append(p) + return groups + + @staticmethod + def apply_1d_nms(group_proposals): + if len(group_proposals) <= 1: + return group_proposals + + unknown_end_proposals = [p for p in group_proposals if p.get("end") == "unknown"] + valid_proposals = [p for p in group_proposals if p.get("end") != "unknown"] + + valid_proposals.sort( + key=lambda x: ( + x.get("confidence_score", 0.0), + -x.get("duration", 0.0), + ), + reverse=True, + ) + + n = len(valid_proposals) + suppressed_mask = [False] * n + + for i in range(n): + if suppressed_mask[i]: + continue + a = (valid_proposals[i]["start"], valid_proposals[i]["end"]) + for j in range(i + 1, n): + if suppressed_mask[j]: + continue + b = (valid_proposals[j]["start"], valid_proposals[j]["end"]) + iou = ConflictResolutionService.calculate_1d_iou(a, b) + iomin = ConflictResolutionService.calculate_1d_iomin(a, b) + + if iou > 0.5 or iomin > 0.8: + suppressed_mask[j] = True + + kept_proposals = [valid_proposals[i] for i in range(n) if not suppressed_mask[i]] + return kept_proposals + unknown_end_proposals + + @classmethod + def clean(cls, blackboard_entries): + try: + sanitized = [cls.sanitize_interval(dict(e)) for e in blackboard_entries] + groups = cls.group_by_semantic_class(sanitized) + + cleaned_entries = [] + for target, proposals in groups.items(): + if len(proposals) <= 1: + cleaned_entries.extend(proposals) + else: + cleaned_entries.extend(cls.apply_1d_nms(proposals)) + + suppressed_count = len(blackboard_entries) - len(cleaned_entries) + if suppressed_count > 0: + logger.info( + "NMS filtering telemetry: Successfully suppressed" + f" {suppressed_count} redundant proposals." + ) + + # Sort chronologically by start time before returning so downstream systems read it in order + cleaned_entries.sort(key=lambda x: float(x.get("start", 0.0))) + + return cleaned_entries + except Exception as e: + logger.error(f"NMS filter crashed: {e}. Falling back to raw ledger state.") + return blackboard_entries diff --git a/artemis/agents/video_analyzer/video_analyzer.md b/artemis/agents/video_analyzer/video_analyzer.md new file mode 100644 index 0000000..215627d --- /dev/null +++ b/artemis/agents/video_analyzer/video_analyzer.md @@ -0,0 +1,51 @@ +# Dynamic Video Analyzer - Main Agent + +You are the Main Agent of the Dynamic Video Analyzer system. Your goal is to process video files of varying lengths and sizes to answer user queries efficiently, avoiding context bloat. + +You receive: +1. **Time Description**: A description of the time range of interest. +2. **Purpose**: The specific goal or question about the video. + +--- + +## Your Capabilities + +You are equipped with specialized tools (`extract_segment_metadata`, `spawn_sub_agent`, `analyze_audio_only`) to extract metadata from video segments and delegate targeted multimodal or audio-only analysis to fault-tolerant sub-agents. + + +--- + +## Execution Logic & Routing Strategy + +### 1. Routing between Video and Audio Analysis +**Probing Metadata**: If you are uncertain about the file size or duration of the requested timeframe, you **should first call `extract_segment_metadata`** to probe the metadata. Use the returned `file_size_mb` and `duration_seconds` to decide whether to proceed with a visual sub-agent, split the request, or fall back to audio-only. + +You must carefully choose the appropriate tool based on the nature of the task: + +- **Must use Visual Analysis (`spawn_sub_agent`)**: + - **UI Interactions**: Identifying clicks, gestures, typing, navigation. + - **Visual State**: Reading on-screen text (OCR), checking UI status (toggles, dialogs, loading states). + - **Visual Content**: Identifying images, layout issues, animations. +- **Can use Audio-Only Analysis (`analyze_audio_only`)**: + - **Speech-centric**: Summarizing spoken content, identifying speakers, transcription. + - **Audio cues**: Detecting specific sounds, music, notification chimes. + - *Constraint*: Only use for long videos where visual info is strictly unnecessary for the query. + +### 2. Formulating Specific Prompts for Sub-Agents +When calling `spawn_sub_agent` or `analyze_audio_only`, the `specific_query` must be clear, detailed, and actionable. +- **Guideline**: Clearly state what to look/listen for and the expected format. Avoid generic instructions. +- **Examples**: + - *Bad*: "Analyze this segment." + - *Good (Audio)*: "Identify the main topic discussed by the speaker and summarize their key conclusions." + - *Good (Video)*: "Observe if the user clicks the 'Confirm' button and verify if the screen transitions to the success page." + +### Dynamic Adaptation & Error Handling +- **Be Flexible**: Do not be rigid about the paths. If you start with Strategy A but find the file size is too large after cropping, switch to Strategy B immediately. +- **Clarification on Ambiguity**: If the `Time Description` is ambiguous or invalid and you cannot reasonably map it to a timeframe, do not guess blindly. Instead, return a clear message to the user explaining that the time specified is unclear and ask for clarification. +- **Minimize Token Usage**: Choose the path that answers the query effectively while minimizing context size. + +--- + +## Final Answer Formatting + +When providing your final answer, you should output a single, direct paragraph of plain text. Include all necessary details. Avoid Markdown formatting, bullet points, headers, or introductory preambles. \ No newline at end of file diff --git a/artemis/agents/video_analyzer/video_analyzer.py b/artemis/agents/video_analyzer/video_analyzer.py new file mode 100644 index 0000000..564956f --- /dev/null +++ b/artemis/agents/video_analyzer/video_analyzer.py @@ -0,0 +1,2347 @@ +# 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 asyncio +import base64 +from datetime import datetime +import glob +import json +import math +import os +from pathlib import Path +import random +import re +import shutil +import tempfile +import time +from typing import Any + +from google import genai +from google.genai import types +from google.genai.errors import APIError +from langchain_core.messages import ( + BaseMessage, + HumanMessage, + SystemMessage, + ToolMessage, +) + +from artemis.agents.video_analyzer.conflict_resolution import ( + ConflictResolutionService, +) +from artemis.config import settings +from artemis.constants import SAFETY_SETTINGS_BLOCK_NONE +from artemis.context import ArtemisContext +from artemis.controllers.controller_factory import get_controller +from artemis.data_engine.trace import CURRENT_TRACE_ID, TraceSpan, trace +from artemis.services.llm import get_llm +from artemis.utils.logger import get_logger +from artemis.utils.video import ( + compress_video_for_api, + extract_audio_from_video, + extract_keyframes_from_video, +) + +try: + from datetime import UTC +except ImportError: + UTC = UTC + +logger = get_logger(__name__) + +TRANSCODE_SEMAPHORE = asyncio.Semaphore(2) +MAIN_AGENT_SEMAPHORE = asyncio.Semaphore(2) +API_SEMAPHORE = asyncio.Semaphore(5) +SLOWDOWN_THRESHOLD_SECONDS = 30.0 +SLOWDOWN_FACTOR = 2.0 +_LAST_CLEANUP_TIME = 0.0 +_CLEANUP_LOCK = asyncio.Lock() + + +UNIVERSAL_SUBMIT_ANSWER_TOOL: dict[str, Any] = { + "type": "function", + "function": { + "name": "submit_answer", + "description": "Submit the final findings for the requested segment analysis.", + "parameters": { + "type": "object", + "properties": { + "summary": { + "type": "string", + "description": "A short paragraph summarizing your findings.", + }, + "analysis": { + "type": "string", + "description": "A short paragraph of detailed analysis.", + }, + "timeline_events": { + "type": "array", + "description": ( + "List of detected events matching the query. If no relevant" + " events occurred, return an empty array." + ), + "items": { + "type": "object", + "properties": { + "start_time": { + "type": "number", + "description": "Event start time (seconds)", + }, + "end_time": { + "type": "number", + "description": "Event end time (seconds)", + }, + "transcription": { + "type": "string", + "description": ( + "A concise transcription of what happened in the event." + ), + }, + "confidence_score": { + "type": "number", + "description": "Confidence score between 0.0 and 1.0", + }, + "verification_timestamp_secs": { + "type": "number", + "description": ( + "Exact timestamp (seconds) for verification screenshot." + ), + }, + }, + "required": [ + "start_time", + "end_time", + "transcription", + "confidence_score", + "verification_timestamp_secs", + ], + }, + }, + }, + "required": ["summary", "analysis"], + }, + }, +} + +UNIVERSAL_MAIN_TOOLS: list[dict[str, Any]] = [ + { + "type": "function", + "function": { + "name": "extract_segment_metadata", + "description": ( + "Retrieves metadata (duration_seconds, file_size_mb) for a video segment." + " Does not return video content or perform content analysis." + ), + "parameters": { + "type": "object", + "properties": { + "start_time": { + "type": "number", + "description": "Start time in seconds", + }, + "end_time": { + "type": "number", + "description": ( + "End time in seconds (optional, defaults to latest available time)" + ), + }, + }, + "required": ["start_time"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "spawn_sub_agent", + "description": ( + "Spawns a sub-agent to analyze visual and audio content of a video segment." + " Parallelizes chunk analysis across workers. Returns combined summary." + ), + "parameters": { + "type": "object", + "properties": { + "start_time": { + "type": "number", + "description": "Start time in seconds", + }, + "end_time": { + "type": "number", + "description": "End time in seconds", + }, + "specific_query": { + "type": "string", + "description": "Specific query or intent for the sub-agent", + }, + }, + "required": ["start_time", "specific_query"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "analyze_audio_only", + "description": ( + "Extracts and analyzes only the audio track from a video segment." + " More efficient for non-visual tasks." + ), + "parameters": { + "type": "object", + "properties": { + "start_time": { + "type": "number", + "description": "Start time in seconds", + }, + "end_time": { + "type": "number", + "description": "End time in seconds (optional)", + }, + "specific_query": { + "type": "string", + "description": "Specific query or intent for the audio content", + }, + }, + "required": ["start_time", "specific_query"], + }, + }, + }, +] + + +async def cleanup_abandoned_gemini_files(client) -> None: + """Scan and delete remaining cloud video files whose creation time has elapsed standard TTL.""" + global _LAST_CLEANUP_TIME + if _CLEANUP_LOCK.locked(): + return + + async with _CLEANUP_LOCK: + now_ts = time.time() + if now_ts - _LAST_CLEANUP_TIME < 3600: + return + _LAST_CLEANUP_TIME = now_ts + + try: + pager = await asyncio.wait_for(client.aio.files.list(), timeout=30) + files = [] + async for f in pager: + files.append(f) + now = datetime.now(UTC) + files_to_delete = [] + + for f in files: + display_name = getattr(f, "display_name", "") or "" + + if ( + display_name.startswith("compressed_") + or display_name.startswith("audio_") + or "artemis" in display_name + ): + create_time = getattr(f, "created_at", None) or getattr(f, "create_time", None) + if not create_time: + continue + + if isinstance(create_time, str): + try: + parsed_time = datetime.fromisoformat(create_time.replace("Z", "+00:00")) + except Exception: + continue + elif isinstance(create_time, datetime): + parsed_time = create_time + else: + continue + + if parsed_time.tzinfo is None: + parsed_time = parsed_time.replace(tzinfo=UTC) + + age_seconds = (now - parsed_time).total_seconds() + + if age_seconds > 7200: + files_to_delete.append(f) + logger.info( + f"Marked expired cloud asset for deletion: {f.name}" + f" ({display_name}), Age: {age_seconds / 3600:.1f}h" + ) + + if files_to_delete: + logger.info(f"Purging {len(files_to_delete)} expired cloud assets in parallel...") + tasks = [ + asyncio.wait_for(client.aio.files.delete(name=f.name), timeout=30) + for f in files_to_delete + ] + await asyncio.gather(*tasks, return_exceptions=True) + except Exception as e: + logger.warning(f"Routine cloud maintenance skipped: {e}") + + +class VideoAnalyzer: + def __init__(self, ctx: ArtemisContext): + self.ctx = ctx + self.blackboard_entries = [] + agent_config = getattr(ctx, "agent_config", None) + self.enable_ledger = getattr(agent_config, "enable_video_ledger", True) + self.local_files_to_cleanup = set() + self.local_dirs_to_cleanup = set() + self.cloud_files_to_cleanup = set() + + sub_prompt_path = Path(__file__).parent / "video_sub_agent.md" + self.sub_system_prompt = ( + sub_prompt_path.read_text(encoding="utf-8") if sub_prompt_path.exists() else "" + ) + + audio_prompt_path = Path(__file__).parent / "video_sub_agent_audio.md" + self.audio_system_prompt = ( + audio_prompt_path.read_text(encoding="utf-8") if audio_prompt_path.exists() else "" + ) + + self._init_tools() + self._init_engine() + + def _init_tools(self) -> None: + """Initializes native and sub-agent tool declarations.""" + self.tools_declaration = [ + types.FunctionDeclaration( + name="extract_segment_metadata", + description=( + "Retrieves metadata (such as duration_seconds and" + " file_size_mb) for a specific video segment. Does not" + " return the video content or perform content analysis." + ), + parameters=types.Schema( + type=types.Type.OBJECT, + properties={ + "start_time": types.Schema( + type=types.Type.NUMBER, + description="Start time in seconds", + ), + "end_time": types.Schema( + type=types.Type.NUMBER, + description=( + "End time in seconds (optional, defaults to latest available time)" + ), + ), + }, + required=["start_time"], + ), + ), + types.FunctionDeclaration( + name="spawn_sub_agent", + description=( + "Spawns a sub-agent to analyze both visual and audio" + " content of a video segment. Parallelizes chunk analysis" + " across workers. Returns combined summary and analysis." + ), + parameters=types.Schema( + type=types.Type.OBJECT, + properties={ + "start_time": types.Schema( + type=types.Type.NUMBER, + description="Start time in seconds", + ), + "end_time": types.Schema( + type=types.Type.NUMBER, + description=( + "End time in seconds. Use the maximum available" + " time if not specified." + ), + ), + "specific_query": types.Schema( + type=types.Type.STRING, + description=("Specific query or intent for the sub-agent"), + ), + }, + required=["start_time", "specific_query"], + ), + ), + types.FunctionDeclaration( + name="analyze_audio_only", + description=( + "Extracts and analyzes only the audio track from a video" + " segment. More efficient for non-visual tasks." + ), + parameters=types.Schema( + type=types.Type.OBJECT, + properties={ + "start_time": types.Schema( + type=types.Type.NUMBER, + description="Start time in seconds", + ), + "end_time": types.Schema( + type=types.Type.NUMBER, + description="End time in seconds (optional)", + ), + "specific_query": types.Schema( + type=types.Type.STRING, + description=("Specific query or intent for the audio content"), + ), + }, + required=["start_time", "specific_query"], + ), + ), + ] + + self.submit_answer_tool = types.FunctionDeclaration( + name="submit_answer", + description=("Submit the final findings for the requested segment analysis."), + parameters=types.Schema( + type=types.Type.OBJECT, + properties={ + "summary": types.Schema( + type=types.Type.STRING, + description=("A short paragraph summarizing of your findings."), + ), + "analysis": types.Schema( + type=types.Type.STRING, + description="A short paragraph of detailed analysis.", + ), + "timeline_events": types.Schema( + type=types.Type.ARRAY, + description=( + "List of detected events matching the query. If no" + " relevant events occurred in the segment, return" + " an empty array." + ), + items=types.Schema( + type=types.Type.OBJECT, + properties={ + "start_time": types.Schema( + type=types.Type.NUMBER, + description="Event start time (seconds)", + ), + "end_time": types.Schema( + type=types.Type.NUMBER, + description="Event end time (seconds)", + ), + "transcription": types.Schema( + type=types.Type.STRING, + description=( + "A concise transcription of what happened in the event. " + ), + ), + "confidence_score": types.Schema( + type=types.Type.NUMBER, + description=("Confidence score between 0 and 1"), + ), + "verification_timestamp_secs": types.Schema( + type=types.Type.NUMBER, + description=( + "Exact timestamp (seconds) for verification screenshot." + ), + ), + }, + required=[ + "start_time", + "end_time", + "transcription", + "confidence_score", + "verification_timestamp_secs", + ], + ), + ), + }, + required=["summary", "analysis", "timeline_events"], + ), + ) + + def _init_engine(self) -> None: + """Initializes model engine and decides whether to use native Gemini or Universal path.""" + ctx = self.ctx + llm_config = getattr(ctx, "llm_config", None) + utils_cfg = getattr(llm_config, "utils", None) if llm_config else None + llm_cfg = getattr(utils_cfg, "video_analyzer", None) if utils_cfg else None + model_str = ( + llm_cfg.model if (llm_cfg and hasattr(llm_cfg, "model")) else "gemini-3.6-flash" + ).lower() + self.model_name = ( + llm_cfg.model if (llm_cfg and hasattr(llm_cfg, "model")) else "gemini-3.6-flash" + ) + if "/" in self.model_name: + self.model_name = self.model_name.split("/")[-1] + + has_google_key = bool( + settings.GOOGLE_API_KEY and settings.GOOGLE_API_KEY.get_secret_value() + ) + is_gemini_model = "gemini" in model_str + + # If client is already set on ctx or explicit Gemini configuration + self.client = getattr(ctx, "_genai_client", None) + if self.client is not None: + self.use_native_gemini = True + elif has_google_key and is_gemini_model: + try: + self.client = genai.Client(api_key=settings.GOOGLE_API_KEY.get_secret_value()) + ctx._genai_client = self.client + self.use_native_gemini = True + asyncio.create_task(cleanup_abandoned_gemini_files(self.client)) + except Exception as e: + logger.warning( + f"Failed to initialize native Gemini client: {e}. Using universal engine." + ) + self.use_native_gemini = False + else: + self.use_native_gemini = False + + @property + def blackboard_ledger(self) -> str: + self.blackboard_entries = ConflictResolutionService.clean(self.blackboard_entries) + if not self.blackboard_entries: + return "No video segments analyzed yet." + lines = [] + for e in self.blackboard_entries: + lines.append(f"{e['start']}s - {e['end']}s: {e['summary']}") + return " | ".join(lines) + + async def upload_and_poll_file(self, compressed_path: Path) -> any: + logger.info(f"Uploading {compressed_path} to Gemini File API...") + file_size_mb = compressed_path.stat().st_size / (1024 * 1024) + upload_timeout = max(30.0, min(120.0, file_size_mb * 2.0)) + + file = await asyncio.wait_for( + self.client.aio.files.upload(file=compressed_path), + timeout=upload_timeout, + ) + self.cloud_files_to_cleanup.add(file.name) + + max_wait = max(60, min(180, 60 + int(file_size_mb * 2))) + + wait_interval = 0.5 + start_wait = time.time() + retry_count = 0 + max_retries = 3 + + while True: + try: + f_state = await asyncio.wait_for( + self.client.aio.files.get(name=file.name), timeout=20 + ) + retry_count = 0 + except Exception as poll_error: + retry_count += 1 + logger.warning( + f"Temporary issue polling file {file.name} (attempt" + f" {retry_count}/{max_retries}): {poll_error}" + ) + if retry_count > max_retries: + raise RuntimeError( + f"Failed to poll Gemini File API after {max_retries} attempts: {poll_error}" + ) + await asyncio.sleep(2.0) + continue + + if f_state.state.name == "ACTIVE": + break + elif f_state.state.name == "FAILED": + raise RuntimeError(f"Gemini File API processing failed for {file.name}") + + if time.time() - start_wait > max_wait: + raise TimeoutError( + f"Gemini File API processing timeout for {file.name}" + f" (Waited {max_wait}s for size {file_size_mb:.1f}MB)" + ) + + logger.info(f"File {file.name} is {f_state.state.name}, waiting {wait_interval}s...") + await asyncio.sleep(wait_interval) + wait_interval = min(3.0, wait_interval * 1.5) + return file + + def get_overlapping_warnings(self, start_time: float, end_time: float | None) -> list[dict]: + warnings = [] + if not isinstance(start_time, (int, float)) or ( + end_time is not None and not isinstance(end_time, (int, float)) + ): + return warnings + if end_time is None or (end_time - start_time) <= 0: + return warnings + req_duration = end_time - start_time + for entry in self.blackboard_entries: + if isinstance(entry, dict) and entry.get("end") != "unknown": + try: + if not isinstance(entry.get("start"), (int, float, str)) or not isinstance( + entry.get("end"), (int, float, str) + ): + continue + overlap_start = max(start_time, float(entry["start"])) + overlap_end = min(end_time, float(entry["end"])) + overlap_duration = max(0.0, overlap_end - overlap_start) + if (overlap_duration / req_duration) > 0.8: + warnings.append( + { + "target": entry.get("target", ""), + "summary": entry.get("summary", ""), + } + ) + except (ValueError, TypeError): + pass + return warnings + + @trace(type="tool", name="extract_segment_metadata") + async def exec_extract_segment_metadata( + self, start_time: float, end_time: float | None = None + ) -> str: + controller = get_controller(self.ctx) + async with TRANSCODE_SEMAPHORE: + result = await controller.extract_segment_metadata(start_time, end_time) + if result.success and result.video_path: + path = Path(result.video_path) + self.local_files_to_cleanup.add(path) + self.local_dirs_to_cleanup.add(path.parent) + size_mb = ( + result.file_size_mb + if result.file_size_mb is not None + else (path.stat().st_size / (1024 * 1024) if path.exists() else 0.0) + ) + + if result.duration_seconds is not None: + duration = result.duration_seconds + elif end_time is not None: + duration = end_time - start_time + else: + try: + duration_cmd = [ + "ffprobe", + "-v", + "error", + "-show_entries", + "format=duration", + "-of", + "default=noprint_wrappers=1:nokey=1", + str(path), + ] + proc = await asyncio.create_subprocess_exec( + *duration_cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + stdout, _ = await proc.communicate() + duration = float(stdout.decode().strip()) + except Exception as e: + logger.warning(f"Failed to resolve duration via ffprobe for {path}: {e}") + duration = 0.0 + if isinstance(duration, (int, float)): + duration = round(duration, 1) + + return json.dumps( + { + "duration_seconds": duration, + "file_size_mb": ( + round(size_mb, 2) if isinstance(size_mb, (int, float)) else size_mb + ), + } + ) + raise RuntimeError(f"Failed to get video segment: {result.message}") + + @trace(type="tool", name="spawn_sub_agent") + async def exec_spawn_sub_agent( + self, start_time: float, end_time: float | None, specific_query: str + ) -> str: + + controller = get_controller(self.ctx) + async with TRANSCODE_SEMAPHORE: + metadata_result = await controller.extract_segment_metadata(start_time, end_time) + if not metadata_result.success: + return f"Error fetching segment metadata: {metadata_result.message}" + + duration = getattr(metadata_result, "duration_seconds", None) + if not isinstance(duration, (int, float)): + duration = ( + (end_time - start_time) + if (isinstance(end_time, (int, float)) and isinstance(start_time, (int, float))) + else 60.0 + ) + if isinstance(duration, (int, float)): + duration = round(duration, 1) + target_chunk_size = 60.0 + num_workers = min(5, math.ceil(duration / target_chunk_size)) + actual_chunk_size = duration / num_workers + chunks = [] + for i in range(num_workers): + cs = round(start_time + (i * actual_chunk_size), 1) + ce = ( + round(cs + actual_chunk_size, 1) + if i < num_workers - 1 + else round(start_time + duration, 1) + ) + chunks.append((cs, ce)) + logger.info( + f"Dynamically chunking video into {num_workers} workers. Chunk" + f" size: {actual_chunk_size:.1f}s" + ) + + async def bounded_spawn(cs, ce, query): + async with API_SEMAPHORE: + return await self._exec_single_chunk(cs, ce, query) + + tasks = [bounded_spawn(cs, ce, specific_query) for cs, ce in chunks] + results = await asyncio.gather(*tasks, return_exceptions=True) + + full_text_logs = [] + for chunk, res in zip(chunks, results): + if isinstance(res, Exception): + full_text_logs.append(f"Chunk {chunk[0]:.1f}s - {chunk[1]:.1f}s FAILED: {res}") + logger.error(f"Chunk failed: {res}") + else: + full_text_logs.append(f"Chunk {chunk[0]:.1f}s - {chunk[1]:.1f}s {res}") + + # Apply NMS to the ledger now that all chunks have appended their events + self.blackboard_entries = ConflictResolutionService.clean(self.blackboard_entries) + + valid_results = [res for res in results if isinstance(res, str)] + if not valid_results: + return "All sub-agent chunks failed." + return " ".join(valid_results) + + async def _exec_single_chunk( + self, start_time: float, end_time: float | None, specific_query: str + ) -> str: + max_retries = 2 + current_start = round(start_time, 1) if isinstance(start_time, (int, float)) else start_time + current_end = round(end_time, 1) if isinstance(end_time, (int, float)) else end_time + + for attempt in range(max_retries + 1): + try: + logger.info( + f"Attempt {attempt + 1}/{max_retries + 1} for sub-agent:" + f" start={current_start}, end={current_end}" + ) + controller = get_controller(self.ctx) + async with TRANSCODE_SEMAPHORE: + result = await controller.extract_segment_metadata(current_start, current_end) + if not result.success or not result.video_path: + raise Exception(f"Failed to get video segment: {result.message}") + + path = Path(result.video_path) + self.local_files_to_cleanup.add(path) + self.local_dirs_to_cleanup.add(path.parent) + + actual_start = getattr(result, "actual_start_relative_time", current_start) + if not isinstance(actual_start, (int, float)): + actual_start = current_start + + slowdown_factor = 1.0 + if ( + current_end is not None + and (current_end - current_start) <= SLOWDOWN_THRESHOLD_SECONDS + ): + slowdown_factor = SLOWDOWN_FACTOR + + async with TRANSCODE_SEMAPHORE: + compressed_path = await compress_video_for_api( + path, + force_compress=True, + start_offset_seconds=actual_start, + slowdown_factor=slowdown_factor, + ) + if compressed_path != path: + self.local_files_to_cleanup.add(compressed_path) + self.local_dirs_to_cleanup.add(compressed_path.parent) + + actual_start = getattr(result, "actual_start_relative_time", current_start) + if not isinstance(actual_start, (int, float)): + actual_start = current_start + + warnings = self.get_overlapping_warnings(start_time, end_time) + warning_block = "" + if warnings: + lines = [ + "WARNING: The following queries were already tried in this timeframe:\n" + ] + for w in warnings: + lines.append(f"Searched for: {w['target']} -> {w['summary']}") + warning_block = "\n".join(lines) + "\n\n" + + duration_secs = getattr(result, "duration_seconds", None) + actual_end = ( + actual_start + duration_secs + if isinstance(duration_secs, (int, float)) + else None + ) + end_str = f" to {actual_end:.1f}s" if actual_end is not None else "" + warning_val = getattr(result, "warning", None) + truncation_note = ( + f" (NOTE: {warning_val})" + if isinstance(warning_val, str) and warning_val + else "" + ) + + slowdown_note = "" + if slowdown_factor != 1.0: + slowdown_note = ( + "WARNING: This video is slowed down to capture fast micro-actions.\n\n" + ) + + prompt_with_context = ( + "IMPORTANT CONTEXT: This video segment corresponds to" + " the test's relative time from" + f" {actual_start:.1f}s{end_str}{truncation_note}. A" + " timestamp showing the exact test relative time" + " (formatted as ' s') is burned in the" + " top-right corner of the video. Please watch this" + " burned-in timestamp and use its value to report all" + " timestamps (e.g., if the burned timestamp shows '90" + f" s', report 90.0s).\n\n{slowdown_note}" + ) + if self.enable_ledger: + prompt_with_context += ( + "BLACKBOARD" + f" LEDGER:\n{self.blackboard_ledger}\nDirective:" + " Use the Blackboard to understand what happened" + " outside your assigned timeframe. Do not" + " duplicate findings already on the board.\n\n" + ) + if warning_block: + prompt_with_context += ( + f"{warning_block}Directive: Avoid repeating" + " searches that have already failed. Try to" + " approach the problem from a different angle or" + " use a different search strategy.\n\n" + ) + prompt_with_context += f"{specific_query}" + + if not self.use_native_gemini: + return await self._exec_single_chunk_universal( + compressed_path=compressed_path, + raw_path=path, + start_time=current_start, + end_time=current_end, + actual_start=actual_start, + prompt_with_context=prompt_with_context, + specific_query=specific_query, + ) + + file = None + try: + with TraceSpan(name="upload_video_to_gemini") as span: + file = await self.upload_and_poll_file(compressed_path) + span.result = f"Uploaded {file.name}" + logger.info( + "Invoking Gemini for sub-agent task with model" + f" {self.model_name} (streaming)..." + ) + trace_id = CURRENT_TRACE_ID.get() + + if self.ctx.data_engine and trace_id: + self.ctx.data_engine.record_trace( + type="llm_call", + name="video_sub_agent", + payload={ + "contents": [ + f"file://{file.name}", + prompt_with_context, + ], + "system_instruction": self.sub_system_prompt, + }, + parent_trace_id=trace_id, + ) + + sub_agent_contents = [ + types.Content( + role="user", + parts=[ + types.Part.from_uri(file_uri=file.uri, mime_type=file.mime_type) + if hasattr(file, "uri") + else file, + types.Part.from_text(text=prompt_with_context), + ], + ) + ] + + sub_max_iterations = 2 + sub_iterations = 0 + final_analysis = "No analysis provided." + final_summary = "No summary provided." + final_full_text = "" + + while sub_iterations < sub_max_iterations: + sub_iterations += 1 + if sub_iterations == sub_max_iterations - 1: + sub_agent_contents.append( + types.Content( + role="user", + parts=[ + types.Part.from_text( + text=( + "[WARNING] You must call the" + " submit_answer tool now, time" + " is running out." + ) + ) + ], + ) + ) + + with TraceSpan(name="gemini_stream_content_main") as span: + stream = await asyncio.wait_for( + self.client.aio.models.generate_content_stream( + model=self.model_name, + contents=sub_agent_contents, + config=types.GenerateContentConfig( + system_instruction=self.sub_system_prompt, + tools=[ + types.Tool( + function_declarations=[self.submit_answer_tool] + ) + ], + safety_settings=SAFETY_SETTINGS_BLOCK_NONE, + ), + ), + timeout=60, + ) + + async def read_stream(target_stream, current_trace_id): + text = "" + function_calls = [] + accumulated_parts = [] + async for chunk in target_stream: + if hasattr(chunk, "usage_metadata") and chunk.usage_metadata: + span.payload["usage_metadata"] = { + "prompt_token_count": getattr( + chunk.usage_metadata, + "prompt_token_count", + 0, + ), + "candidates_token_count": getattr( + chunk.usage_metadata, + "candidates_token_count", + 0, + ), + "total_token_count": getattr( + chunk.usage_metadata, + "total_token_count", + 0, + ), + } + + chunk_text = chunk.text or "" + if chunk_text: + text += chunk_text + if self.ctx.data_engine and current_trace_id: + self.ctx.data_engine.stream_output( + current_trace_id, chunk_text + ) + if chunk.function_calls: + function_calls.extend(chunk.function_calls) + + candidates = getattr(chunk, "candidates", None) + if ( + candidates + and isinstance(candidates, list) + and len(candidates) > 0 + ): + content = getattr(candidates[0], "content", None) + parts = getattr(content, "parts", None) if content else None + if parts: + for part in parts: + accumulated_parts.append(part) + if getattr(part, "thought", False) and part.text: + if self.ctx.data_engine and current_trace_id: + self.ctx.data_engine.stream_output( + current_trace_id, + part.text, + is_thinking=True, + ) + else: + if chunk.function_calls: + for fc in chunk.function_calls: + accumulated_parts.append(types.Part(function_call=fc)) + + return text, function_calls, accumulated_parts + + full_text, function_calls, original_parts = await asyncio.wait_for( + read_stream(stream, trace_id), timeout=180 + ) + final_full_text += full_text + "\n" + + if function_calls: + answered = False + for fc in function_calls: + if ( + fc.name.split(":")[-1] if ":" in fc.name else fc.name + ) == "submit_answer": + args = fc.args + timeline_events = args.get("timeline_events", []) + final_summary = ( + args.get("summary", "No summary provided.") + .strip() + .replace("\n", " ") + ) + final_analysis = ( + args.get("analysis", "No analysis provided.") + .strip() + .replace("\n", " ") + ) + + for event in timeline_events: + if ( + "confidence_score" not in event + or not isinstance( + event["confidence_score"], + (int, float), + ) + or isinstance(event["confidence_score"], bool) + or not (0.0 <= float(event["confidence_score"]) <= 1.0) + ): + error_msg = ( + "Error: 'confidence_score' is" + " missing or out of bounds in" + " one of the timeline_events." + " It must be a float between" + " 0.0 and 1.0 inclusive. Fix" + " the payload and call" + " submit_answer again." + ) + sub_agent_contents.append( + types.Content( + role="model", + parts=original_parts + if original_parts + else [ + types.Part(function_call=fc) + for fc in function_calls + ] + + ( + [types.Part.from_text(text=full_text)] + if full_text + else [] + ), + ) + ) + sub_agent_contents.append( + types.Content( + role="user", + parts=[ + types.Part.from_function_response( + name=fc.name, + response={"error": (error_msg)}, + ) + ], + ) + ) + answered = "error" + break + + if ( + "verification_timestamp_secs" not in event + or not isinstance( + event["verification_timestamp_secs"], + (int, float), + ) + or isinstance( + event["verification_timestamp_secs"], bool + ) + ): + error_msg = ( + "Error:" + " 'verification_timestamp_secs'" + " is missing or invalid in one" + " of the timeline_events. Fix" + " the payload and call" + " submit_answer again." + ) + sub_agent_contents.append( + types.Content( + role="model", + parts=original_parts + if original_parts + else [ + types.Part(function_call=fc) + for fc in function_calls + ] + + ( + [types.Part.from_text(text=full_text)] + if full_text + else [] + ), + ) + ) + sub_agent_contents.append( + types.Content( + role="user", + parts=[ + types.Part.from_function_response( + name=fc.name, + response={"error": (error_msg)}, + ) + ], + ) + ) + answered = "error" + break + + verification_ts = float( + event.get( + "verification_timestamp_secs", + start_time, + ) + ) + confidence = float(event.get("confidence_score", 0.0)) + + if verification_ts < start_time or ( + end_time is not None and verification_ts > end_time + ): + error_msg = ( + "Error:" + " 'verification_timestamp_secs'" + f" {verification_ts} must fall" + " within the start_time" + f" ({start_time}) and end_time" + f" ({end_time}) boundaries. Fix" + " the payload and call" + " submit_answer again." + ) + sub_agent_contents.append( + types.Content( + role="model", + parts=original_parts + if original_parts + else [ + types.Part(function_call=fc) + for fc in function_calls + ] + + ( + [types.Part.from_text(text=full_text)] + if full_text + else [] + ), + ) + ) + sub_agent_contents.append( + types.Content( + role="user", + parts=[ + types.Part.from_function_response( + name=fc.name, + response={"error": (error_msg)}, + ) + ], + ) + ) + answered = "error" + break + + if answered == "error": + break + + for event in timeline_events: + verification_ts = float( + event.get( + "verification_timestamp_secs", + start_time, + ) + ) + confidence = float(event.get("confidence_score", 0.0)) + relative_offset = max(0.0, verification_ts - actual_start) + + out_dir = Path("/tmp/video_analyzer") + out_dir.mkdir(parents=True, exist_ok=True) + existing_imgs = glob.glob(f"{out_dir}/img_*.jpg") + max_idx = -1 + for img in existing_imgs: + match = re.search(r"img_(\d+)\.jpg$", img) + if match: + max_idx = max(max_idx, int(match.group(1))) + next_idx = max_idx + 1 + screenshot_output_path = out_dir / f"img_{next_idx}.jpg" + while screenshot_output_path.exists(): + next_idx += 1 + screenshot_output_path = out_dir / f"img_{next_idx}.jpg" + screenshot_output_path.touch() + self.local_files_to_cleanup.add(screenshot_output_path) + ffmpeg_cmd = [ + "ffmpeg", + "-ss", + str(relative_offset), + "-i", + str(path), + "-vframes", + "1", + "-q:v", + "2", + "-y", + str(screenshot_output_path), + ] + try: + proc = await asyncio.create_subprocess_exec( + *ffmpeg_cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + await proc.communicate() + except Exception as e: + logger.warning(f"Failed to extract frame: {e}") + + screenshot_path_str = str(screenshot_output_path) + entry = { + "start": event.get("start_time", start_time), + "end": event.get("end_time", end_time), + "target": specific_query, + "summary": event.get("transcription", ""), + "confidence_score": confidence, + } + if Path(screenshot_output_path).exists(): + entry["screenshot"] = str(screenshot_path_str) + + self.blackboard_entries.append(entry) + + answered = True + break + + if answered is True: + break + elif answered == "error": + pass + else: + sub_agent_contents.append( + types.Content( + role="model", + parts=original_parts + if original_parts + else [types.Part(function_call=fc) for fc in function_calls] + + ( + [types.Part.from_text(text=full_text)] + if full_text + else [] + ), + ) + ) + sub_agent_contents.append( + types.Content( + role="user", + parts=[ + types.Part.from_function_response( + name=function_calls[0].name, + response={ + "error": ( + "Tool not recognized." + " Please use" + " submit_answer." + ) + }, + ) + ], + ) + ) + else: + sub_agent_contents.append( + types.Content( + role="model", + parts=[types.Part.from_text(text=full_text)], + ) + ) + + end_str = f"{current_end:.1f}s" if current_end is not None else "unknown" + return ( + f"[from {current_start:.1f}s to {end_str}] Summary:" + f" {final_summary} Analysis: {final_analysis}" + ) + finally: + if file: + logger.info(f"Cleaning up cloud file {file.name}...") + try: + await asyncio.wait_for( + self.client.aio.files.delete(name=file.name), + timeout=30, + ) + self.cloud_files_to_cleanup.discard(file.name) + except Exception as ce: + logger.error(f"Failed to delete cloud file {file.name}: {ce}") + + except APIError as api_err: + logger.warning(f"Sub-agent APIError attempt {attempt + 1} failed: {api_err}") + code = getattr(api_err, "code", None) + if code == 429: + if attempt >= max_retries: + raise api_err + + backoff = (2.0**attempt) + random.uniform(0.1, 1.0) + logger.info( + f"Rate limited (429). Backing off for {backoff:.2f}s" + " without expanding boundaries..." + ) + await asyncio.sleep(backoff) + continue + elif code in [400, 401, 403, 404]: + raise api_err + + if attempt >= max_retries: + raise api_err + backoff = 2.0**attempt + await asyncio.sleep(backoff) + except Exception as e: + logger.warning(f"Sub-agent attempt {attempt + 1} failed: {e}") + if attempt >= max_retries: + raise e + + backoff = 2.0**attempt + logger.info(f"Backing off for {backoff}s before retry...") + await asyncio.sleep(backoff) + + @trace(type="tool", name="analyze_audio_only") + async def exec_analyze_audio_only( + self, start_time: float, end_time: float | None, specific_query: str + ) -> str: + max_retries = 2 + current_start = round(start_time, 1) if isinstance(start_time, (int, float)) else start_time + current_end = round(end_time, 1) if isinstance(end_time, (int, float)) else end_time + + for attempt in range(max_retries + 1): + try: + logger.info( + f"Attempt {attempt + 1}/{max_retries + 1} for" + f" analyze_audio_only: start={current_start}," + f" end={current_end}" + ) + controller = get_controller(self.ctx) + async with TRANSCODE_SEMAPHORE: + result = await controller.extract_segment_metadata(current_start, current_end) + if not result.success or not result.video_path: + raise Exception(f"Failed to get video segment: {result.message}") + + path = Path(result.video_path) + self.local_files_to_cleanup.add(path) + self.local_dirs_to_cleanup.add(path.parent) + + async with TRANSCODE_SEMAPHORE: + audio_path = await extract_audio_from_video(path) + self.local_files_to_cleanup.add(audio_path) + self.local_dirs_to_cleanup.add(audio_path.parent) + actual_start = getattr(result, "actual_start_relative_time", current_start) + if not isinstance(actual_start, (int, float)): + actual_start = current_start + + warnings = self.get_overlapping_warnings(start_time, end_time) + warning_block = "" + if warnings: + lines = [ + "WARNING: The following queries were already tried in this timeframe:\n" + ] + for w in warnings: + lines.append(f"Searched for: {w['target']} -> {w['summary']}") + warning_block = "\n".join(lines) + "\n\n" + + duration_secs = getattr(result, "duration_seconds", None) + actual_end = ( + actual_start + duration_secs + if isinstance(duration_secs, (int, float)) + else None + ) + end_str = f" to {actual_end:.1f}s" if actual_end is not None else "" + warning_val = getattr(result, "warning", None) + truncation_note = ( + f" (NOTE: {warning_val})" + if isinstance(warning_val, str) and warning_val + else "" + ) + + prompt_with_context = ( + "IMPORTANT CONTEXT: This audio segment corresponds to" + " the test's relative time from" + f" {actual_start:.1f}s{end_str}{truncation_note}. When" + " you report timestamps, remember that your 00:00 is" + f" exactly {actual_start:.1f}s in the test relative" + " time. Please calculate and output all timestamps as" + " the exact test relative time (e.g., if you hear" + " something at 00:05, report" + f" {actual_start + 5:.1f}s).\n\n" + ) + if self.enable_ledger: + prompt_with_context += ( + "BLACKBOARD" + f" LEDGER:\n{self.blackboard_ledger}\nDirective:" + " Use the Blackboard to understand what happened" + " outside your assigned timeframe. Do not" + " duplicate findings already on the board.\n\n" + ) + if warning_block: + prompt_with_context += ( + f"{warning_block}Directive: Do not repeat the" + " failed searches listed above. Shift your" + " attention to the current task.\n\n" + ) + prompt_with_context += f"{specific_query}" + + if not self.use_native_gemini: + return await self._exec_analyze_audio_universal( + audio_path=audio_path, + start_time=current_start, + end_time=current_end, + actual_start=actual_start, + prompt_with_context=prompt_with_context, + specific_query=specific_query, + ) + + file = None + try: + with TraceSpan(name="upload_audio_to_gemini") as span: + file = await self.upload_and_poll_file(audio_path) + span.result = f"Uploaded {file.name}" + logger.info( + f"Invoking Gemini for audio-only task with model {self.model_name}..." + ) + + sub_agent_contents = [ + types.Content( + role="user", + parts=[ + types.Part.from_uri(file_uri=file.uri, mime_type=file.mime_type) + if hasattr(file, "uri") + else file, + types.Part.from_text(text=prompt_with_context), + ], + ) + ] + + sub_max_iterations = 2 + sub_iterations = 0 + final_confidence_score = 0.0 + final_summary = "No summary provided." + final_analysis = "No analysis provided." + final_full_text = "" + + while sub_iterations < sub_max_iterations: + sub_iterations += 1 + if sub_iterations == sub_max_iterations - 1: + sub_agent_contents.append( + types.Content( + role="user", + parts=[ + types.Part.from_text( + text=( + "[WARNING] You must call the" + " submit_answer tool now, time" + " is running out." + ) + ) + ], + ) + ) + + with TraceSpan(name="gemini_generate_audio_content") as span: + response = await asyncio.wait_for( + self.client.aio.models.generate_content( + model=self.model_name, + contents=sub_agent_contents, + config=types.GenerateContentConfig( + system_instruction=self.audio_system_prompt, + tools=[ + types.Tool( + function_declarations=[self.submit_answer_tool] + ) + ], + safety_settings=SAFETY_SETTINGS_BLOCK_NONE, + ), + ), + timeout=180, + ) + + full_text = response.text or "" + final_full_text += full_text + "\n" + function_calls = response.function_calls or [] + + if function_calls: + answered = False + for fc in function_calls: + if ( + fc.name.split(":")[-1] if ":" in fc.name else fc.name + ) == "submit_answer": + args = fc.args + if ( + "confidence_score" not in args + or not isinstance( + args["confidence_score"], + (int, float), + ) + or not (0.0 <= float(args["confidence_score"]) <= 1.0) + ): + error_msg = ( + "Error: 'confidence_score' is" + " missing or out of bounds. It must" + " be a float between 0.0 and 1.0" + " inclusive. Fix the payload and" + " call submit_answer again." + ) + sub_agent_contents.append( + types.Content( + role="model", + parts=response.candidates[0].content.parts, + ) + ) + sub_agent_contents.append( + types.Content( + role="user", + parts=[ + types.Part.from_function_response( + name=fc.name, + response={"error": error_msg}, + ) + ], + ) + ) + answered = "error" + break + + final_summary = ( + args.get("summary", "No summary provided.") + .strip() + .replace("\n", " ") + ) + final_analysis = ( + args.get("analysis", "No analysis provided.") + .strip() + .replace("\n", " ") + ) + final_confidence_score = float(args["confidence_score"]) + answered = True + break + + if answered is True: + break + elif answered == "error": + pass + else: + sub_agent_contents.append( + types.Content( + role="model", + parts=response.candidates[0].content.parts, + ) + ) + sub_agent_contents.append( + types.Content( + role="user", + parts=[ + types.Part.from_function_response( + name=function_calls[0].name, + response={ + "error": ( + "Tool not recognized." + " Please use" + " submit_answer." + ) + }, + ) + ], + ) + ) + else: + sub_agent_contents.append( + types.Content( + role="model", + parts=[types.Part.from_text(text=full_text)], + ) + ) + + end_val = end_time if end_time is not None else "unknown" + self.blackboard_entries.append( + { + "start": start_time, + "end": end_val, + "target": specific_query, + "summary": final_summary, + "confidence_score": final_confidence_score, + } + ) + + end_str = f"{end_time:.1f}s" if end_time is not None else "unknown" + return ( + f"[from {start_time:.1f}s to {end_str}] Summary:" + f" {final_summary} Analysis: {final_analysis}" + ) + finally: + if file: + logger.info(f"Cleaning up cloud file {file.name}...") + try: + await asyncio.wait_for( + self.client.aio.files.delete(name=file.name), + timeout=30, + ) + self.cloud_files_to_cleanup.discard(file.name) + except Exception as ce: + logger.error(f"Failed to delete cloud file {file.name}: {ce}") + + except APIError as api_err: + logger.warning(f"Audio APIError attempt {attempt + 1} failed: {api_err}") + code = getattr(api_err, "code", None) + if code == 429: + if attempt >= max_retries: + raise api_err + + backoff = (2.0**attempt) + random.uniform(0.1, 1.0) + logger.info( + f"Rate limited (429). Backing off for {backoff:.2f}s" + " without expanding boundaries..." + ) + await asyncio.sleep(backoff) + continue + elif code in [400, 401, 403, 404]: + raise api_err + + if attempt >= max_retries: + raise api_err + backoff = 2.0**attempt + await asyncio.sleep(backoff) + except Exception as e: + logger.warning(f"Audio attempt {attempt + 1} failed: {e}") + if attempt >= max_retries: + raise e + + backoff = 2.0**attempt + logger.info(f"Backing off for {backoff}s before retry...") + await asyncio.sleep(backoff) + + async def _exec_single_chunk_universal( + self, + compressed_path: Path, + raw_path: Path, + start_time: float, + end_time: float | None, + actual_start: float, + prompt_with_context: str, + specific_query: str, + ) -> str: + """Executes sub-agent chunk analysis via keyframe extraction and LangChain ChatModel.""" + keyframes = extract_keyframes_from_video( + compressed_path, + fps=1.0, + max_frames=30, + max_dimension=1080, + ) + + user_blocks: list[dict[str, Any]] = [{"type": "text", "text": prompt_with_context}] + + for ts_sec, frame_bytes in keyframes: + b64_str = base64.b64encode(frame_bytes).decode("utf-8") + abs_time = actual_start + ts_sec + user_blocks.append( + { + "type": "text", + "text": f"--- Video Keyframe at {abs_time:.1f}s ---", + } + ) + user_blocks.append( + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{b64_str}"}, + } + ) + + messages = [ + SystemMessage(content=self.sub_system_prompt), + HumanMessage(content=user_blocks), + ] + + trace_id = CURRENT_TRACE_ID.get() + if self.ctx.data_engine and trace_id: + self.ctx.data_engine.record_trace( + type="llm_call", + name="video_sub_agent_universal", + payload={ + "num_keyframes": len(keyframes), + "prompt": prompt_with_context, + "system_instruction": self.sub_system_prompt, + }, + parent_trace_id=trace_id, + ) + + llm = get_llm(self.ctx, name="video_analyzer") + bound_llm = llm.bind_tools([UNIVERSAL_SUBMIT_ANSWER_TOOL]) + + with TraceSpan(name="universal_sub_agent_call") as span: + response = await asyncio.wait_for(bound_llm.ainvoke(messages), timeout=180) + span.result = ( + f"Received response (tool_calls={len(getattr(response, 'tool_calls', []))})" + ) + + final_summary = "No summary provided." + final_analysis = "No analysis provided." + timeline_events = [] + + if getattr(response, "tool_calls", None): + for tc in response.tool_calls: + if tc.get("name") == "submit_answer": + args = tc.get("args", {}) + final_summary = ( + str(args.get("summary", "")).strip().replace("\n", " ") or final_summary + ) + final_analysis = ( + str(args.get("analysis", "")).strip().replace("\n", " ") or final_analysis + ) + timeline_events = args.get("timeline_events", []) + break + elif getattr(response, "content", None): + text_content = str(response.content).strip() + final_summary = text_content[:200].replace("\n", " ") + final_analysis = text_content + + # Process timeline events & save screenshots for the blackboard ledger + for event in timeline_events: + if not isinstance(event, dict): + continue + v_ts = float(event.get("verification_timestamp_secs", start_time)) + conf = float(event.get("confidence_score", 0.8)) + + screenshot_path_str = None + try: + rel_offset = max(0.0, v_ts - actual_start) + out_dir = Path("/tmp/video_analyzer") + out_dir.mkdir(parents=True, exist_ok=True) + + existing_imgs = glob.glob(f"{out_dir}/img_*.jpg") + max_idx = -1 + for img in existing_imgs: + m = re.search(r"img_(\d+)\.jpg$", img) + if m: + max_idx = max(max_idx, int(m.group(1))) + next_idx = max_idx + 1 + shot_path = out_dir / f"img_{next_idx}.jpg" + + ffmpeg_cmd = [ + "ffmpeg", + "-ss", + str(rel_offset), + "-i", + str(raw_path), + "-vframes", + "1", + "-q:v", + "2", + "-y", + str(shot_path), + ] + proc = await asyncio.create_subprocess_exec( + *ffmpeg_cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + await proc.communicate() + if shot_path.exists(): + self.local_files_to_cleanup.add(shot_path) + screenshot_path_str = str(shot_path) + except Exception as fe: + logger.warning(f"Failed extracting keyframe for event {v_ts}s: {fe}") + + entry = { + "start": event.get("start_time", start_time), + "end": event.get("end_time", end_time), + "target": specific_query, + "summary": event.get("transcription", final_summary), + "confidence_score": conf, + } + if screenshot_path_str: + entry["screenshot"] = screenshot_path_str + self.blackboard_entries.append(entry) + + end_str = f"{end_time:.1f}s" if end_time is not None else "unknown" + return ( + f"[from {start_time:.1f}s to {end_str}] Summary: {final_summary}" + f" Analysis: {final_analysis}" + ) + + async def _exec_analyze_audio_universal( + self, + audio_path: Path, + start_time: float, + end_time: float | None, + actual_start: float, + prompt_with_context: str, + specific_query: str, + ) -> str: + """Executes audio analysis via Universal LangChain ChatModel.""" + messages = [ + SystemMessage(content=self.audio_system_prompt), + HumanMessage(content=prompt_with_context), + ] + llm = get_llm(self.ctx, name="video_analyzer") + bound_llm = llm.bind_tools([UNIVERSAL_SUBMIT_ANSWER_TOOL]) + + response = await asyncio.wait_for(bound_llm.ainvoke(messages), timeout=180) + final_summary = "No audio summary provided." + final_analysis = "No audio analysis provided." + + if getattr(response, "tool_calls", None): + for tc in response.tool_calls: + if tc.get("name") == "submit_answer": + args = tc.get("args", {}) + final_summary = ( + str(args.get("summary", "")).strip().replace("\n", " ") or final_summary + ) + final_analysis = ( + str(args.get("analysis", "")).strip().replace("\n", " ") or final_analysis + ) + break + elif getattr(response, "content", None): + final_summary = str(response.content)[:200].replace("\n", " ") + final_analysis = str(response.content) + + end_val = end_time if end_time is not None else "unknown" + self.blackboard_entries.append( + { + "start": start_time, + "end": end_val, + "target": specific_query, + "summary": final_summary, + "confidence_score": 0.8, + } + ) + + end_str = f"{end_time:.1f}s" if end_time is not None else "unknown" + return ( + f"[from {start_time:.1f}s to {end_str}] Audio Summary: {final_summary}" + f" Analysis: {final_analysis}" + ) + + async def _run_universal( + self, time_description: str, purpose: str, system_prompt: str + ) -> tuple[str, str]: + """Runs VideoAnalyzer main reasoning loop using Universal LangChain BaseChatModel.""" + llm = get_llm(self.ctx, name="video_analyzer") + bound_llm = llm.bind_tools(UNIVERSAL_MAIN_TOOLS) + + messages: list[BaseMessage] = [ + SystemMessage(content=system_prompt), + HumanMessage( + content=( + f"Time description: {time_description}\n" + f"Purpose: {purpose}\n" + "Please coordinate video metadata extraction and sub-agent analysis." + ) + ), + ] + + max_iterations = 6 + iterations = 0 + agent_outcome = "" + status = "success" + + try: + while iterations < max_iterations: + iterations += 1 + response = await asyncio.wait_for(bound_llm.ainvoke(messages), timeout=300) + messages.append(response) + + if not getattr(response, "tool_calls", None): + agent_outcome = str(response.content) + break + + for tc in response.tool_calls: + t_name = tc["name"] + t_args = tc.get("args", {}) + t_id = tc.get("id", f"call_{iterations}") + logger.info(f"Executing universal tool '{t_name}' with args {t_args}") + + try: + if t_name == "extract_segment_metadata": + res = await self.exec_extract_segment_metadata( + t_args.get("start_time"), t_args.get("end_time") + ) + elif t_name == "spawn_sub_agent": + res = await self.exec_spawn_sub_agent( + t_args.get("start_time"), + t_args.get("end_time"), + t_args.get("specific_query"), + ) + elif t_name == "analyze_audio_only": + res = await self.exec_analyze_audio_only( + t_args.get("start_time"), + t_args.get("end_time"), + t_args.get("specific_query"), + ) + else: + res = f"Error: Tool '{t_name}' is not recognized." + except Exception as tool_err: + logger.error(f"Error executing tool '{t_name}': {tool_err}") + res = f"Error executing tool '{t_name}': {tool_err}" + + messages.append(ToolMessage(tool_call_id=t_id, name=t_name, content=str(res))) + + if not agent_outcome: + agent_outcome = "Video analyzer completed all turns without error." + + except Exception as e: + logger.error(f"Universal video analyzer failed: {e}") + agent_outcome = f"Video analysis failed: {e}" + status = "error" + + return agent_outcome, status + + @trace(type="agent", name="video_analyzer") + async def run(self, time_description: str, purpose: str) -> tuple[str, str]: + ctx = self.ctx + async with MAIN_AGENT_SEMAPHORE: + logger.info( + f"Video analyzer invoked with time: '{time_description}', purpose: '{purpose}'" + ) + + # Initialize engine & cleanup tracking + self._init_engine() + + # Resolve model name and temperature from config + llm_config = getattr(ctx, "llm_config", None) + utils_cfg = getattr(llm_config, "utils", None) if llm_config else None + llm_cfg = getattr(utils_cfg, "video_analyzer", None) if utils_cfg else None + temperature = 0.2 + thinking_level = None + if llm_cfg: + self.model_name = llm_cfg.model + if "/" in self.model_name: + self.model_name = self.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 + else: + self.model_name = "gemini-3.6-flash" + + # Track files for cleanup + self.local_files_to_cleanup = set() + self.local_dirs_to_cleanup = set() + self.cloud_files_to_cleanup = set() + + # 1. Use initialized native tools declarations + tools_declaration = self.tools_declaration + + sub_prompt_path = Path(__file__).parent / "video_sub_agent.md" + self.sub_system_prompt = sub_prompt_path.read_text(encoding="utf-8") + + audio_prompt_path = Path(__file__).parent / "video_sub_agent_audio.md" + self.audio_system_prompt = audio_prompt_path.read_text(encoding="utf-8") + + prompt_path = Path(__file__).parent / "video_analyzer.md" + system_prompt = prompt_path.read_text(encoding="utf-8") + if self.enable_ledger: + system_prompt += ( + "\n## Trust Protocol\nConsult Blackboard Ledger before" + " searching. Evaluate if new search is necessary based on" + " previous context. Explain reasoning before aborting or" + " proceeding.\n" + ) + + if not self.use_native_gemini: + return await self._run_universal( + time_description=time_description, + purpose=purpose, + system_prompt=system_prompt, + ) + + contents = [ + types.Content( + role="user", + parts=[ + types.Part.from_text( + text=(f"Time description: {time_description}\nPurpose: {purpose}") + ) + ], + ) + ] + + agent_outcome = "" + status = "success" + max_iterations = 10 + iterations = 0 + last_ledger_index = 0 + + try: + while iterations < max_iterations: + iterations += 1 + logger.info(f"Iteration {iterations}: Invoking Native Gemini SDK...") + + self.blackboard_entries = ConflictResolutionService.clean(self.blackboard_entries) + + new_entries = self.blackboard_entries[last_ledger_index:] + if self.enable_ledger and new_entries: + lines = [] + image_parts = [] + for e in new_entries: + line = f"[{e['start']}s - {e['end']}s] SUMMARY: {e['summary']}" + screenshot_file = e.get("screenshot") or e.get("screenshot_path") + if screenshot_file and Path(screenshot_file).exists(): + line += f" PROOF: {Path(screenshot_file).name}" + image_parts.append( + types.Part.from_text(text=f"PROOF: {Path(screenshot_file).name}") + ) + image_parts.append( + types.Part.from_bytes( + data=Path(screenshot_file).read_bytes(), + mime_type="image/jpeg", + ) + ) + lines.append(line) + diff_text = "NEW BLACKBOARD ENTRIES ADDED: " + " | ".join(lines) + contents.append( + types.Content( + role="user", + parts=[types.Part.from_text(text=diff_text)], + ) + ) + if image_parts: + contents.append(types.Content(role="user", parts=image_parts)) + last_ledger_index = len(self.blackboard_entries) + + is_final_turn = iterations == max_iterations + if is_final_turn: + contents.append( + types.Content( + role="user", + parts=[ + types.Part.from_text( + text=( + "This is your final iteration; all" + " tools are stripped, and you must" + " provide your final answer directly." + ) + ) + ], + ) + ) + else: + if iterations > 1: + contents.append( + types.Content( + role="user", + parts=[ + types.Part.from_text( + text=( + "You have not completed the video" + " analysis yet (iteration" + f" {iterations} of" + f" {max_iterations})." + ) + ) + ], + ) + ) + + with TraceSpan(name="gemini_main_agent_call") as span: + trace_id = CURRENT_TRACE_ID.get() + + if ctx.data_engine and trace_id: + serialized_contents = [] + for c in contents: + if hasattr(c, "model_dump"): + dumped = c.model_dump() + if ( + isinstance(dumped, dict) + and "parts" in dumped + and dumped["parts"] + ): + for p in dumped["parts"]: + if ( + isinstance(p, dict) + and "inline_data" in p + and p["inline_data"] + and "data" in p["inline_data"] + ): + p["inline_data"]["data"] = "" + serialized_contents.append(dumped) + else: + serialized_contents.append(str(c)) + + ctx.data_engine.record_trace( + type="llm_call", + name="video_analyzer_main", + payload={ + "contents": serialized_contents, + "system_instruction": system_prompt, + "tools": ( + [t.name for t in tools_declaration] + if (tools_declaration and not is_final_turn) + else [] + ), + }, + parent_trace_id=trace_id, + ) + + async def run_stream(): + stream = await asyncio.wait_for( + self.client.aio.models.generate_content_stream( + model=self.model_name, + contents=contents, + config=types.GenerateContentConfig( + system_instruction=system_prompt, + temperature=temperature, + tools=[] + if is_final_turn + else [types.Tool(function_declarations=tools_declaration)], + safety_settings=SAFETY_SETTINGS_BLOCK_NONE, + **( + { + "thinking_config": types.ThinkingConfig( + thinking_level=thinking_level + ) + } + if thinking_level + else {} + ), + ), + ), + timeout=60, + ) + + text = "" + accumulated_parts = [] + function_calls = [] + + async for chunk in stream: + if hasattr(chunk, "usage_metadata") and chunk.usage_metadata: + span.payload["usage_metadata"] = { + "prompt_token_count": getattr( + chunk.usage_metadata, + "prompt_token_count", + 0, + ), + "candidates_token_count": getattr( + chunk.usage_metadata, + "candidates_token_count", + 0, + ), + "total_token_count": getattr( + chunk.usage_metadata, + "total_token_count", + 0, + ), + } + + chunk_text = chunk.text or "" + text += chunk_text + if ctx.data_engine and trace_id and chunk_text: + ctx.data_engine.stream_output( + trace_id, chunk_text, is_thinking=False + ) + + if chunk.function_calls: + function_calls.extend(chunk.function_calls) + + candidates = getattr(chunk, "candidates", None) + if candidates and isinstance(candidates, list) and len(candidates) > 0: + content = getattr(candidates[0], "content", None) + parts = getattr(content, "parts", None) if content else None + if parts: + for part in parts: + accumulated_parts.append(part) + if getattr(part, "thought", False) and part.text: + if ctx.data_engine and trace_id: + ctx.data_engine.stream_output( + trace_id, + part.text, + is_thinking=True, + ) + else: + if chunk.function_calls: + for fc in chunk.function_calls: + accumulated_parts.append(types.Part(function_call=fc)) + + class DummyResponse: + def __init__(self, text, function_calls, parts): + self.text = text + self.function_calls = function_calls + self.candidates = [ + types.Candidate( + content=types.Content( + role="model", + parts=parts + if parts + else ( + [types.Part.from_text(text=text)] if text else [] + ), + ) + ) + ] + + return DummyResponse(text, function_calls, accumulated_parts) + + response = await asyncio.wait_for(run_stream(), timeout=300) + span.result = ( + f"Function calls: {len(response.function_calls)}" + if response.function_calls + else "Final answer" + ) + + function_calls = response.function_calls + if not function_calls: + agent_outcome = response.text + break + + if response.candidates and response.candidates[0].content: + contents.append(response.candidates[0].content) + else: + contents.append( + types.Content( + role="model", + parts=[types.Part(function_call=fc) for fc in function_calls], + ) + ) + + sub_agent_calls = [ + fc + for fc in function_calls + if (fc.name.split(":")[-1] if ":" in fc.name else fc.name) == "spawn_sub_agent" + ] + audio_only_calls = [ + fc + for fc in function_calls + if (fc.name.split(":")[-1] if ":" in fc.name else fc.name) + == "analyze_audio_only" + ] + other_calls = [ + fc + for fc in function_calls + if (fc.name.split(":")[-1] if ":" in fc.name else fc.name) + not in ["spawn_sub_agent", "analyze_audio_only"] + ] + + tool_response_parts = [] + + if sub_agent_calls: + logger.info(f"Executing {len(sub_agent_calls)} sub-agent calls in parallel...") + + async def bounded_spawn(fc): + async with API_SEMAPHORE: + return await self.exec_spawn_sub_agent( + fc.args.get("start_time"), + fc.args.get("end_time"), + fc.args.get("specific_query"), + ) + + tasks = [bounded_spawn(fc) for fc in sub_agent_calls] + results = await asyncio.gather(*tasks, return_exceptions=True) + + for fc, result in zip(sub_agent_calls, results): + if isinstance(result, Exception): + error_msg = f"Sub-agent failed. Error: {result}" + logger.error(error_msg) + tool_response_parts.append( + types.Part.from_function_response( + name=fc.name, response={"error": error_msg} + ) + ) + else: + tool_response_parts.append( + types.Part.from_function_response( + name=fc.name, response={"result": result} + ) + ) + + if audio_only_calls: + logger.info( + f"Executing {len(audio_only_calls)} analyze_audio_only calls in parallel..." + ) + + async def bounded_audio_spawn(fc): + async with API_SEMAPHORE: + return await self.exec_analyze_audio_only( + fc.args.get("start_time"), + fc.args.get("end_time"), + fc.args.get("specific_query"), + ) + + tasks = [bounded_audio_spawn(fc) for fc in audio_only_calls] + results = await asyncio.gather(*tasks, return_exceptions=True) + + for fc, result in zip(audio_only_calls, results): + if isinstance(result, Exception): + error_msg = f"Audio analysis failed. Error: {result}" + logger.error(error_msg) + tool_response_parts.append( + types.Part.from_function_response( + name=fc.name, response={"error": error_msg} + ) + ) + else: + tool_response_parts.append( + types.Part.from_function_response( + name=fc.name, response={"result": result} + ) + ) + + for fc in other_calls: + name = fc.name.split(":")[-1] if ":" in fc.name else fc.name + args = fc.args + logger.info(f"Executing tool '{name}' sequentially...") + + if name == "extract_segment_metadata": + try: + res = await self.exec_extract_segment_metadata( + args.get("start_time"), args.get("end_time") + ) + tool_response_parts.append( + types.Part.from_function_response( + name=name, response={"result": res} + ) + ) + except Exception as e: + tool_response_parts.append( + types.Part.from_function_response( + name=name, response={"error": str(e)} + ) + ) + + if tool_response_parts: + contents.append(types.Content(role="user", parts=tool_response_parts)) + + if iterations >= max_iterations and not agent_outcome: + agent_outcome = ( + "Error: Video analyzer agent reached maximum iterations" + " without settling on a final answer." + ) + status = "error" + + except Exception as e: + logger.error(f"Video analyzer failed: {e}") + agent_outcome = f"Video analysis failed: {e}" + status = "error" + finally: + if self.cloud_files_to_cleanup: + logger.info( + f"Cleaning up {len(self.cloud_files_to_cleanup)} remaining" + " cloud files in parallel..." + ) + tasks = [ + asyncio.wait_for(self.client.aio.files.delete(name=name), timeout=30) + for name in list(self.cloud_files_to_cleanup) + ] + await asyncio.gather(*tasks, return_exceptions=True) + self.cloud_files_to_cleanup.clear() + + if os.environ.get("KEEP_VIDEOS") or os.environ.get("ARTEMIS_DEBUG"): + logger.info( + "Skipping cleanup of local video files and directories due to debug mode." + ) + else: + for path in list(self.local_files_to_cleanup): + try: + if path.exists(): + logger.info(f"Cleaning up local file {path}...") + path.unlink() + except Exception as e: + logger.error(f"Failed to delete local file {path}: {e}") + + for dir_path in list(self.local_dirs_to_cleanup): + try: + resolved = dir_path.resolve() + system_dirs = { + Path("/").resolve(), + Path("/tmp").resolve(), + Path("/var/tmp").resolve(), + Path.home().resolve(), + Path(tempfile.gettempdir()).resolve(), + } + if resolved in system_dirs: + logger.debug(f"Skipping deletion of root/system directory: {dir_path}") + continue + + if dir_path.exists(): + logger.info(f"Cleaning up local directory {dir_path}...") + shutil.rmtree(dir_path, ignore_errors=True) + except Exception as e: + logger.error(f"Failed to delete local directory {dir_path}: {e}") + + return agent_outcome, status diff --git a/artemis/agents/video_analyzer/video_sub_agent.md b/artemis/agents/video_analyzer/video_sub_agent.md new file mode 100644 index 0000000..c25bfdb --- /dev/null +++ b/artemis/agents/video_analyzer/video_sub_agent.md @@ -0,0 +1,15 @@ +# Video Sub-Agent - Chunk Analyzer + +You are a specialized video chunk analyzer. Your task is to analyze a specific segment of a video and answer the query provided by the Main Agent. + +You must follow these principles strictly: +1. **Causality and Chronology**: First establish a chronology of states or events on screen over the entire video segment to help differentiate between user inputs (e.g., physical taps, touch indicators) and system state changes. +2. **Objective Observations**: Describe only what you see in the video. Pay extreme attention to mobile UI specifics: identify which icons are clicked (visual touch indicators), screen transitions, loading spinners, popup dialogs, keyboard appearing/disappearing, and any changes in UI text or status bars. Use fuzzy matching when evaluating UI text and notifications. Do not assume or speculate about things not visible in this chunk. +3. **Context Isolation**: You only see a part of the full video. Do not assume what happened before or after this segment. +4. **Focus on Intent**: Answer the specific query or prompt provided by the Main Agent. Do not give a generic summary if a specific question is asked. +5. **Structured Output**: Provide your findings clearly, using the exact global timestamps calculated according to the IMPORTANT CONTEXT prepended to your prompt. You must rely on the timestamp burned directly into the video frames to accurately report event timings. Never use relative 00:xx formats; always output the final global test time in seconds. +6. **Visual Proof**: Identify the exact best frame that provides visual proof of your findings. Supply this timestamp as `verification_timestamp_secs` when calling `submit_answer`. +7. **Slow-Motion Awareness**: If a slowdown warning is present in prompt context, animations, transitions, and timers appear in slow-motion. Do not misclassify slow-motion as lag, stutter, or frozen UI. + +## Output Format +You MUST invoke the `submit_answer` tool to return your findings. diff --git a/artemis/agents/video_analyzer/video_sub_agent_audio.md b/artemis/agents/video_analyzer/video_sub_agent_audio.md new file mode 100644 index 0000000..f4fa2f1 --- /dev/null +++ b/artemis/agents/video_analyzer/video_sub_agent_audio.md @@ -0,0 +1,12 @@ +# Audio Sub-Agent - Chunk Analyzer + +You are a specialized audio chunk analyzer. Your task is to analyze a specific audio segment extracted from a video and answer the query provided by the Main Agent. + +You must follow these principles strictly: +1. **Objective Observations**: Describe only what you hear in the audio (speech, conversations, background sounds, tone, language). Pay attention not only to the spoken words but also to non-verbal cues (tone, hesitation, emotion) and critical background sounds (notification chimes, click sounds, error beeps) if they are relevant to the query. Do not assume or speculate about things not audible. +2. **Context Isolation**: You only hear a part of the full audio. Do not assume what happened before or after this segment. +3. **Focus on Intent**: Answer the specific query or prompt provided by the Main Agent. Do not give a generic summary if a specific question is asked. +4. **Structured Output**: Provide your findings clearly, using the exact global timestamps calculated according to the IMPORTANT CONTEXT prepended to your prompt. Never use relative 00:xx formats; always output the final global test time in seconds. + +## Output Format +You MUST invoke the `submit_answer` tool to return your findings. diff --git a/artemis/clients/adb_tunnel.py b/artemis/clients/adb_tunnel.py new file mode 100644 index 0000000..2cd79b6 --- /dev/null +++ b/artemis/clients/adb_tunnel.py @@ -0,0 +1,327 @@ +# 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. + +"""ADB WebSocket tunnel implementation. + +Bridges local TCP connections to a remote ADB server via WebSocket. +This replaces the need for the `lim` CLI for ADB tunneling. + +The tunnel runs in a dedicated thread with its own event loop to ensure +it cannot be blocked by synchronous operations in the main thread. + +Based on the Go SDK implementation which: +1. Creates a TCP listener +2. Accepts TCP connections from ADB +3. Dials the WebSocket +4. Forwards traffic bidirectionally +""" + +import asyncio +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager +import socket +import threading + +from artemis.utils.logger import get_logger +import websockets +from websockets.asyncio.client import ClientConnection + +logger = get_logger(__name__) + +BUFFER_SIZE = 32 * 1024 # 32KB, same as Go SDK +PING_INTERVAL = 30 # seconds + + +class AdbTunnel: + """ADB tunnel that bridges local TCP connections to a remote WebSocket endpoint. + + Runs in a dedicated thread with its own event loop to ensure the tunnel + cannot be blocked by synchronous operations in the caller's thread. + + Usage: + tunnel = AdbTunnel(remote_url, token) + addr = await tunnel.start() # Returns "127.0.0.1:PORT" + # Use addr with ADB: adb connect 127.0.0.1:PORT + await tunnel.stop() + """ + + def __init__(self, remote_url: str, token: str): + """Initialize the ADB tunnel. + + Args: + remote_url: WebSocket URL for the remote ADB endpoint + token: Bearer token for authentication + """ + self.remote_url = remote_url + self.token = token + + self._local_port: int | None = None + self._thread: threading.Thread | None = None + self._loop: asyncio.AbstractEventLoop | None = None + self._started = threading.Event() + self._stop_event: asyncio.Event | None = None + + @property + def local_addr(self) -> str: + """Get the local address to connect ADB to.""" + if self._local_port is None: + raise RuntimeError("Tunnel not started") + return f"127.0.0.1:{self._local_port}" + + @property + def local_port(self) -> int: + """Get the local port number.""" + if self._local_port is None: + raise RuntimeError("Tunnel not started") + return self._local_port + + async def start(self) -> str: + """Start the ADB tunnel in a dedicated thread. + + Returns: + The local address (host:port) to connect ADB to. + """ + if self._thread is not None and self._thread.is_alive(): + return self.local_addr + + # Create listener in main thread to get the port immediately + listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + listener.bind(("127.0.0.1", 0)) + listener.listen(5) + listener.setblocking(False) + self._local_port = listener.getsockname()[1] + + logger.info(f"ADB tunnel listening on {self.local_addr}") + + # Start tunnel thread + self._thread = threading.Thread( + target=self._run_in_thread, + args=(listener,), + name=f"adb-tunnel-{self._local_port}", + daemon=True, + ) + self._thread.start() + + # Wait for thread to initialize its event loop + self._started.wait(timeout=5.0) + if not self._started.is_set(): + await self.stop() + raise RuntimeError("Tunnel thread failed to start") + + return self.local_addr + + async def stop(self) -> None: + """Stop the ADB tunnel and cleanup resources.""" + if self._loop and self._stop_event: + # Signal the tunnel loop to stop + self._loop.call_soon_threadsafe(self._stop_event.set) + + if self._thread and self._thread.is_alive(): + await asyncio.to_thread(self._thread.join, 5.0) + if self._thread.is_alive(): + logger.warning("ADB tunnel thread did not stop within 5s") + + self._thread = None + self._loop = None + self._stop_event = None + logger.info("ADB tunnel stopped") + + def _run_in_thread(self, listener: socket.socket) -> None: + """Run the tunnel event loop in a dedicated thread.""" + try: + self._loop = asyncio.new_event_loop() + asyncio.set_event_loop(self._loop) + self._stop_event = asyncio.Event() + self._started.set() + + self._loop.run_until_complete(self._run_tunnel(listener)) + except Exception as e: + logger.error(f"Tunnel thread error: {e}") + finally: + listener.close() + if self._loop: + self._loop.close() + + async def _run_tunnel(self, listener: socket.socket) -> None: + """Run the tunnel - accept TCP connections and bridge to WebSocket.""" + loop = asyncio.get_event_loop() + connection_tasks: list[asyncio.Task] = [] + + stop_event = self._stop_event + assert stop_event is not None + + logger.debug("Tunnel: ready to accept connections") + + while not stop_event.is_set(): + try: + # Use wait_for to periodically check stop_event + try: + tcp_conn, addr = await asyncio.wait_for( + loop.sock_accept(listener), + timeout=1.0, + ) + except TimeoutError: + # Check if we should stop + continue + + logger.debug(f"Tunnel: new TCP connection from {addr}") + + task = asyncio.create_task( + self._handle_connection(tcp_conn, loop), + name=f"conn-{addr[1]}", + ) + task.add_done_callback(lambda t: t.exception() if not t.cancelled() else None) + connection_tasks.append(task) + connection_tasks = [t for t in connection_tasks if not t.done()] + + except Exception as e: + if not stop_event.is_set(): + logger.error(f"Tunnel accept error: {e}") + break + + # Cleanup pending connections + for task in connection_tasks: + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + async def _handle_connection( + self, + tcp_conn: socket.socket, + loop: asyncio.AbstractEventLoop, + ) -> None: + """Handle a single TCP connection.""" + tcp_conn.setblocking(False) + + ws: ClientConnection | None = None + try: + logger.debug(f"Connecting to WebSocket: {self.remote_url[:80]}...") + ws = await websockets.connect( + self.remote_url, + additional_headers={"Authorization": f"Bearer {self.token}"}, + ping_interval=PING_INTERVAL, + ping_timeout=10, + ) + logger.debug("WebSocket connected to remote ADB") + await self._bridge(tcp_conn, ws, loop) + except Exception as e: + logger.error(f"Connection error: {e}") + finally: + if ws is not None: + await ws.close() + tcp_conn.close() + logger.debug("Connection closed") + + async def _bridge( + self, + tcp_conn: socket.socket, + ws: ClientConnection, + loop: asyncio.AbstractEventLoop, + ) -> None: + """Bridge TCP and WebSocket bidirectionally.""" + + stop_event = self._stop_event + assert stop_event is not None + + async def tcp_to_ws() -> None: + """Forward data from TCP to WebSocket.""" + msg_count = 0 + try: + while not stop_event.is_set(): + try: + data = await asyncio.wait_for( + loop.sock_recv(tcp_conn, BUFFER_SIZE), + timeout=1.0, + ) + except TimeoutError: + continue + except (OSError, ConnectionError) as e: + logger.debug(f"tcp->ws: recv error after {msg_count} msgs: {e}") + break + if not data: + logger.debug(f"tcp->ws: closed by client after {msg_count} msgs") + break + msg_count += 1 + logger.debug(f"tcp->ws: [{msg_count}] sending {len(data)} bytes") + await ws.send(data) + logger.debug(f"tcp->ws: loop ended after {msg_count} messages") + except asyncio.CancelledError: + logger.debug(f"tcp->ws: cancelled after {msg_count} messages") + raise + except Exception as e: + logger.error(f"tcp->ws: error after {msg_count} messages: {e}") + + async def ws_to_tcp() -> None: + """Forward data from WebSocket to TCP.""" + msg_count = 0 + try: + logger.debug("ws->tcp: starting message loop") + async for message in ws: + if stop_event.is_set(): + break + if isinstance(message, bytes) and message: + msg_count += 1 + logger.debug(f"ws->tcp: [{msg_count}] received {len(message)} bytes") + await loop.sock_sendall(tcp_conn, message) + elif isinstance(message, str): + logger.debug(f"ws->tcp: received string: {message[:100]}") + logger.debug(f"ws->tcp: loop ended after {msg_count} messages") + except asyncio.CancelledError: + logger.debug(f"ws->tcp: cancelled after {msg_count} messages") + raise + except Exception as e: + logger.error(f"ws->tcp: error after {msg_count} messages: {e}") + + tcp_task = asyncio.create_task(tcp_to_ws()) + ws_task = asyncio.create_task(ws_to_tcp()) + + done, pending = await asyncio.wait( + [tcp_task, ws_task], + return_when=asyncio.FIRST_COMPLETED, + ) + + for task in pending: + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + +@asynccontextmanager +async def adb_tunnel(remote_url: str, token: str) -> AsyncGenerator[AdbTunnel, None]: + """Context manager for creating and managing an ADB tunnel. + + Args: + remote_url: WebSocket URL for the remote ADB endpoint + token: Bearer token for authentication + + Yields: + AdbTunnel instance with the tunnel running + + Example: + async with adb_tunnel(ws_url, token) as tunnel: + # Connect ADB to tunnel.local_addr + adb_client = AdbClient() + adb_client.connect(tunnel.local_addr) + """ + tunnel = AdbTunnel(remote_url, token) + try: + await tunnel.start() + yield tunnel + finally: + await tunnel.stop() diff --git a/artemis/clients/ui_automator_client.py b/artemis/clients/ui_automator_client.py new file mode 100644 index 0000000..6d92507 --- /dev/null +++ b/artemis/clients/ui_automator_client.py @@ -0,0 +1,416 @@ +# 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. + +"""UIAutomator2 client for Android device screen data retrieval. + +Provides an alternative to the Maestro-based screen API with direct device +access. +Handles Maestro blocker detection and removal before connecting. +""" + +import base64 +from io import BytesIO +import re as _re +import subprocess +import time +from typing import TYPE_CHECKING +import xml.etree.ElementTree as ET + +from PIL import Image +from pydantic import BaseModel +import uiautomator2 as u2 + +from artemis.utils.logger import get_logger + +if TYPE_CHECKING: + from uiautomator2 import Device + +logger = get_logger(__name__) + + +MAESTRO_PACKAGE = "dev.mobile.maestro" + + +class _CyFunctionDetectorMeta(type): + def __instancecheck__(self, instance): + name = type(instance).__name__ + return ( + name + in ( + "cyfunction", + "cython_function_or_method", + "builtin_function_or_method", + ) + or "cyfunction" in name.lower() + ) + + +class CyFunctionDetector(metaclass=_CyFunctionDetectorMeta): + pass + + +class UIAutomatorScreenData(BaseModel): + """Screen data response from UIAutomator2.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + + base64: str + hierarchy_xml: str + elements: list[dict] + width: int + height: int + + +def _parse_hierarchy_xml_to_elements(hierarchy_xml: str) -> list[dict]: + """Parse uiautomator2 XML hierarchy into a flat list of element dictionaries. + + The output format matches the existing screen API format with attributes + like: + - resource-id + - text + - content-desc + - bounds (as string "[x1,y1][x2,y2]") + - class + - package + - checkable, checked, clickable, enabled, focusable, focused + - scrollable, long-clickable, password, selected + + Args: + hierarchy_xml: XML string from uiautomator2.dump_hierarchy() + + Returns: + Flat list of element dictionaries + """ + elements: list[dict] = [] + + try: + root = ET.fromstring(hierarchy_xml) + except ET.ParseError as e: + logger.error(f"Failed to parse hierarchy XML: {e}") + return elements + + def _extract_element(node: ET.Element) -> None: + """Recursively extract elements from XML nodes.""" + element: dict = {} + + for attr_name, attr_value in node.attrib.items(): + if attr_name == "resource-id": + element["resource-id"] = attr_value + elif attr_name == "text": + element["text"] = attr_value + elif attr_name == "content-desc": + element["content-desc"] = attr_value + element["accessibilityText"] = attr_value + elif attr_name == "bounds": + element["bounds"] = attr_value + + match = _re.match(r"\[(\d+),(\d+)\]\[(\d+),(\d+)\]", attr_value) + if match: + x1, y1, x2, y2 = map(int, match.groups()) + element["parsed_bounds"] = { + "left": x1, + "top": y1, + "right": x2, + "bottom": y2, + } + elif attr_name == "class": + element["class"] = attr_value + elif attr_name == "package": + element["package"] = attr_value + elif attr_name in ( + "checkable", + "checked", + "clickable", + "enabled", + "focusable", + "focused", + "scrollable", + "long-clickable", + "password", + "selected", + ): + element[attr_name] = attr_value + else: + element[attr_name] = attr_value + + if element: + elements.append(element) + + for child in node: + _extract_element(child) + + _extract_element(root) + + return elements + + +def _pil_to_base64(img: Image.Image, format: str = "JPEG", quality: int = 80) -> str: + """Convert PIL Image to base64 string with high-quality compression.""" + buffer = BytesIO() + if format.upper() in ("JPEG", "JPG"): + format = "JPEG" + if img.mode in ("RGBA", "P"): + img = img.convert("RGB") + img.save(buffer, format=format, quality=quality, optimize=True) + else: + img.save(buffer, format=format) + return base64.b64encode(buffer.getvalue()).decode("utf-8") + + +def _is_package_installed(device_id: str, pkg: str) -> bool: + """Check if a package is installed on the device.""" + try: + result = subprocess.run( + ["adb", "-s", device_id, "shell", "pm", "list", "packages"], + capture_output=True, + text=True, + timeout=10, + ) + if result.returncode != 0: + logger.warning(f"Failed to list packages: {result.stderr}") + return False + + lines = result.stdout.splitlines() + target = f"package:{pkg}" + return target in lines + except subprocess.TimeoutExpired: + logger.warning("Timeout checking installed packages") + return False + except Exception as e: + logger.warning(f"Error checking installed packages: {e}") + return False + + +def _uninstall_package(device_id: str, pkg: str) -> bool: + """Uninstall a package from the device for the current user.""" + try: + result = subprocess.run( + [ + "adb", + "-s", + device_id, + "shell", + "pm", + "uninstall", + "--user", + "0", + pkg, + ], + capture_output=True, + text=True, + timeout=30, + ) + if result.returncode == 0: + logger.info(f"Successfully uninstalled {pkg}") + return True + else: + logger.warning(f"Failed to uninstall {pkg}: {result.stderr}") + return False + except subprocess.TimeoutExpired: + logger.warning(f"Timeout uninstalling {pkg}") + return False + except Exception as e: + logger.warning(f"Error uninstalling {pkg}: {e}") + return False + + +def _ensure_maestro_not_installed(device_id: str) -> None: + """Check if Maestro is installed and uninstall it. + + Maestro conflicts with uiautomator2 - if Maestro's package is installed, + u2.connect() will fail. This function ensures Maestro is removed before + attempting to connect. + """ + if _is_package_installed(device_id, MAESTRO_PACKAGE): + logger.warning( + f"Maestro ({MAESTRO_PACKAGE}) detected - uninstalling to enable UIAutomator2..." + ) + _uninstall_package(device_id, MAESTRO_PACKAGE) + + +class UIAutomatorClient: + """UIAutomator2 client for Android screen data retrieval. + + This client uses uiautomator2 library for direct device communication, + providing faster hierarchy and screenshot retrieval compared to Maestro. + + Important: Maestro must not be installed on the device as it conflicts + with uiautomator2. This client automatically handles Maestro removal. + """ + + def __init__(self, device_id: str): + """Initialize the UIAutomator client. + + Args: + device_id: The Android device serial number (e.g., "emulator-5554") + """ + self._device_id = device_id + self._device: Device | None = None + + def _ensure_connected(self) -> "Device": + """Ensure connection to the device, handling Maestro blocker. + + Returns: + Connected uiautomator2 Device instance + """ + if self._device is not None: + try: + # Quick check if connection is still alive + self._device.info + return self._device + except Exception: + logger.warning("UIAutomator2 connection lost, reconnecting...") + self._device = None + + # Ensure Maestro is not blocking us + _ensure_maestro_not_installed(self._device_id) + + # Connect to device + logger.info(f"Connecting UIAutomator2 to device: {self._device_id}") + self._device = u2.connect(self._device_id) + logger.info("UIAutomator2 connected successfully") + return self._device + + def press_key(self, key: str): + """Press a key on the device. + + Args: + key: Key to press (e.g., "home", "back", "enter"...) + """ + device = self._ensure_connected() + return device.press(key=key) + + def send_text(self, text: str) -> None: + """Send text input to the device using FastInputIME. + + This method supports special characters (e.g., 'รถ') that ADB shell + input text cannot handle. The original IME is restored after input. + + Args: + text: The text to input + """ + + device = self._ensure_connected() + device.set_fastinput_ime(True) + # Allow IME to initialize + time.sleep(0.3) + try: + device.send_keys(text) + # Give FastInputIME time to process the broadcast and commit text + # before switching it off and killing it. + time.sleep(0.5) + finally: + device.set_fastinput_ime(False) + + def clear_text(self) -> bool: + """Clear text of the currently focused element using UIAutomator2. + + Returns True on success, False on failure. + """ + try: + device = self._ensure_connected() + device(focused=True).clear_text() + return True + except Exception as e: + logger.error(f"UIAutomator2 clear_text failed: {e}") + return False + + def get_hierarchy(self) -> str: + """Get the UI hierarchy XML from the device. + + Returns: + Compressed UI hierarchy XML string + """ + device = self._ensure_connected() + return device.dump_hierarchy(compressed=True) + + def get_screenshot(self) -> Image.Image | None: + """Capture a screenshot from the device. + + Uses adb screencap for more reliable capture of all layers (including + popups), falling back to UIAutomator2 if it fails. + + Returns: + PIL Image or None if capture failed + """ + try: + result = subprocess.run( + ["adb", "-s", self._device_id, "exec-out", "screencap", "-p"], + capture_output=True, + timeout=10, + check=True, + ) + return Image.open(BytesIO(result.stdout)) + except Exception as e: + logger.warning( + f"Failed to capture screenshot via adb: {e}, falling back to uiautomator2" + ) + device = self._ensure_connected() + return device.screenshot() + + def get_screenshot_base64(self) -> str | None: + """Capture a screenshot and return as base64 string. + + Returns: + Base64 encoded PNG screenshot or None if capture failed + """ + screenshot = self.get_screenshot() + if screenshot is None: + return None + return _pil_to_base64(screenshot, format="JPEG") + + def get_screen_data(self) -> UIAutomatorScreenData: + """Get complete screen data including screenshot and hierarchy. + + Returns: + UIAutomatorScreenData with screenshot, hierarchy, elements, and + dimensions + """ + device = self._ensure_connected() + + # Get screenshot + screenshot = self.get_screenshot() + if screenshot is None: + raise RuntimeError("Failed to capture screenshot") + + # Get hierarchy + hierarchy_xml = device.dump_hierarchy(compressed=True) + + # Parse XML to flat elements list + elements = _parse_hierarchy_xml_to_elements(hierarchy_xml) + + return UIAutomatorScreenData( + base64=_pil_to_base64(screenshot, format="JPEG"), + hierarchy_xml=hierarchy_xml, + elements=elements, + width=screenshot.width, + height=screenshot.height, + ) + + def disconnect(self) -> None: + """Disconnect from the device.""" + self._device = None + logger.info("UIAutomator2 client disconnected") + + +def get_client(device_id: str) -> UIAutomatorClient: + """Factory function to create a UIAutomatorClient. + + Args: + device_id: The Android device serial number + + Returns: + UIAutomatorClient instance + """ + return UIAutomatorClient(device_id=device_id) diff --git a/artemis/config/__init__.py b/artemis/config/__init__.py new file mode 100644 index 0000000..c7e7616 --- /dev/null +++ b/artemis/config/__init__.py @@ -0,0 +1,253 @@ +# 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. + +"""Unified configuration package and single entry point for ARTEMIS settings, paths, and models.""" + +from artemis.config.agent import ( + AgentGlobalConfig, + CheckerConfig, + CommitteeConfig, + ExplorerConfig, + FlashProfileConfig, + OutputterConfig, + PlannerValidationConfig, + ProProfileConfig, + StepSummarizerConfig, + VideoAnalyzerConfig, + load_agent_config, + resolve_explorer_version, +) +from artemis.config.constants import ( + AGENT_CONFIG_FILENAME, + DATA_ENGINE_DB_FILENAME, + DEFAULT_ACTION_TIMEOUT, + DEFAULT_ADB_HOST, + DEFAULT_ADB_PORT, + DEFAULT_EDGE_MONITOR_PORT, + DEFAULT_EXPLORER_VERSION, + DEFAULT_GATEWAY_PORT, + DEFAULT_MODEL, + DEFAULT_PROFILE, + DEFAULT_STREAM_PING_INTERVAL, + DOTENV_FILENAME, + ENV_ADB_DEVICE_SERIAL, + ENV_ADB_HOST, + ENV_ADB_PORT, + ENV_ADB_SERVER_SOCKET, + ENV_ANTHROPIC_API_KEY, + ENV_ANTIGRAVITY_APP_DIR, + ENV_ANTIGRAVITY_LS_ADDRESS, + ENV_DATA_ENGINE_DB_PATH, + ENV_EVENTS_OUTPUT_PATH, + ENV_GCP_API_KEY, + ENV_GEMINI_API_KEY, + ENV_GOOGLE_API_KEY, + ENV_KEEP_VIDEOS, + ENV_ARTEMIS_APP_DIR, + ENV_ARTEMIS_CLOUD_GATEWAY_URL, + ENV_ARTEMIS_CLOUD_MODE, + ENV_ARTEMIS_CLOUD_SESSION_ID, + ENV_ARTEMIS_DEBUG, + ENV_ARTEMIS_DEFAULT_MODEL, + ENV_ARTEMIS_DEFAULT_PROFILE, + ENV_ARTEMIS_DEVICE_ID, + ENV_ARTEMIS_EDGE_PORT, + ENV_ARTEMIS_EXPLORER_VERSION, + ENV_ARTEMIS_IPC_PORT, + ENV_ARTEMIS_MCP_SERVER, + ENV_ARTEMIS_TENANT_ID, + ENV_ARTEMIS_TENANT_TOKEN, + ENV_ARTEMIS_TRACES_DIR, + ENV_ARTEMIS_USE_FILE_API, + ENV_ARTEMIS_USE_USER_DIR, + ENV_OPEN_ROUTER_API_KEY, + ENV_OPENAI_API_KEY, + ENV_OPENAI_BASE_URL, + ENV_RESULTS_OUTPUT_PATH, + ENV_XAI_API_KEY, + IPC_PORT_FILENAME, + LLM_CONFIG_FILENAME, + LLM_CONFIG_OVERRIDE_FILENAME, + LS_ADDRESS_FILENAME, + AgentNode, + AgentNodeWithFallback, + ExplorerVersion, + LLMProvider, + LLMUtilsNode, + LLMUtilsNodeWithFallback, +) +from artemis.config.llm import ( + CyFunctionDetector, + LLM, + LLMConfig, + LLMConfigUtils, + LLMWithFallback, + deep_merge_llm_config, + get_default_llm_config, + initialize_llm_config, + load_llm_config_override, + parse_llm_config, + validate_vertex_ai_credentials, +) +from artemis.config.output import ( + OutputConfig, + prepare_output_files, + record_events, +) +from artemis.config.paths import ( + CONFIG_DIR, + GLOBAL_APP_DIR, + GLOBAL_JETSKI_DIR, + ROOT_DIR, + get_app_dir, + get_config_path, + get_data_engine_db_path, + get_default_traces_path, + get_ipc_port_file, + get_ls_address_file, + get_temp_dir, + get_traces_dir, + is_frozen_bundle, +) +from artemis.config.runtime import ( + cleanup_temp_dir, + clear_ipc_port, + clear_ls_address, + read_ipc_port, + read_ls_address, + write_ipc_port, + write_ls_address, +) +from artemis.config.settings import ( + Settings, + settings, +) + +__all__ = [ + # Settings & Singleton + "Settings", + "settings", + # Paths & Dirs + "ROOT_DIR", + "CONFIG_DIR", + "GLOBAL_APP_DIR", + "GLOBAL_JETSKI_DIR", + "get_app_dir", + "get_traces_dir", + "get_default_traces_path", + "get_temp_dir", + "get_config_path", + "get_data_engine_db_path", + "get_ipc_port_file", + "get_ls_address_file", + "is_frozen_bundle", + # LLM Config & Models + "CyFunctionDetector", + "LLM", + "LLMWithFallback", + "LLMConfigUtils", + "LLMConfig", + "parse_llm_config", + "initialize_llm_config", + "get_default_llm_config", + "deep_merge_llm_config", + "load_llm_config_override", + "validate_vertex_ai_credentials", + # Agent & Tool Config + "AgentGlobalConfig", + "CheckerConfig", + "CommitteeConfig", + "ExplorerConfig", + "FlashProfileConfig", + "OutputterConfig", + "PlannerValidationConfig", + "ProProfileConfig", + "StepSummarizerConfig", + "VideoAnalyzerConfig", + "load_agent_config", + "resolve_explorer_version", + # Output Config + "OutputConfig", + "prepare_output_files", + "record_events", + # Runtime & IPC State + "read_ipc_port", + "write_ipc_port", + "clear_ipc_port", + "read_ls_address", + "write_ls_address", + "clear_ls_address", + "cleanup_temp_dir", + # Types & Literals + "LLMProvider", + "ExplorerVersion", + "LLMUtilsNode", + "LLMUtilsNodeWithFallback", + "AgentNode", + "AgentNodeWithFallback", + # Filenames + "LLM_CONFIG_FILENAME", + "LLM_CONFIG_OVERRIDE_FILENAME", + "AGENT_CONFIG_FILENAME", + "DATA_ENGINE_DB_FILENAME", + "IPC_PORT_FILENAME", + "LS_ADDRESS_FILENAME", + "DOTENV_FILENAME", + # Defaults + "DEFAULT_ADB_HOST", + "DEFAULT_ADB_PORT", + "DEFAULT_GATEWAY_PORT", + "DEFAULT_EDGE_MONITOR_PORT", + "DEFAULT_ACTION_TIMEOUT", + "DEFAULT_STREAM_PING_INTERVAL", + "DEFAULT_PROFILE", + "DEFAULT_MODEL", + "DEFAULT_EXPLORER_VERSION", + # Environment Variables + "ENV_ARTEMIS_APP_DIR", + "ENV_ANTIGRAVITY_APP_DIR", + "ENV_ARTEMIS_USE_USER_DIR", + "ENV_ARTEMIS_TRACES_DIR", + "ENV_DATA_ENGINE_DB_PATH", + "ENV_ARTEMIS_CLOUD_MODE", + "ENV_ARTEMIS_CLOUD_SESSION_ID", + "ENV_ARTEMIS_CLOUD_GATEWAY_URL", + "ENV_ARTEMIS_TENANT_ID", + "ENV_ARTEMIS_TENANT_TOKEN", + "ENV_ARTEMIS_EDGE_PORT", + "ENV_ARTEMIS_IPC_PORT", + "ENV_ANTIGRAVITY_LS_ADDRESS", + "ENV_ARTEMIS_MCP_SERVER", + "ENV_ADB_DEVICE_SERIAL", + "ENV_ARTEMIS_DEVICE_ID", + "ENV_ADB_HOST", + "ENV_ADB_PORT", + "ENV_ADB_SERVER_SOCKET", + "ENV_GOOGLE_API_KEY", + "ENV_GEMINI_API_KEY", + "ENV_GCP_API_KEY", + "ENV_OPENAI_API_KEY", + "ENV_OPENAI_BASE_URL", + "ENV_ANTHROPIC_API_KEY", + "ENV_OPEN_ROUTER_API_KEY", + "ENV_XAI_API_KEY", + "ENV_ARTEMIS_EXPLORER_VERSION", + "ENV_ARTEMIS_DEFAULT_PROFILE", + "ENV_ARTEMIS_DEFAULT_MODEL", + "ENV_ARTEMIS_USE_FILE_API", + "ENV_KEEP_VIDEOS", + "ENV_ARTEMIS_DEBUG", + "ENV_EVENTS_OUTPUT_PATH", + "ENV_RESULTS_OUTPUT_PATH", +] diff --git a/artemis/config/agent.py b/artemis/config/agent.py new file mode 100644 index 0000000..453b519 --- /dev/null +++ b/artemis/config/agent.py @@ -0,0 +1,490 @@ +# 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. + +"""Agent tools, sub-agent capabilities, and explorer settings configuration.""" + +import os +from pathlib import Path +from typing import Any + +from pydantic import BaseModel, Field, model_validator + +from artemis.config.constants import ( + AGENT_CONFIG_FILENAME, + DEFAULT_EXPLORER_VERSION, + ExplorerVersion, +) +from artemis.config.paths import ROOT_DIR, get_config_path +from artemis.utils.file import load_jsonc +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +class VideoAnalyzerConfig(BaseModel): + """Configuration specific to VideoAnalyzer and screen recording.""" + + enabled: bool | None = Field( + default=None, + description=( + "Whether continuous screen video recording and video analyzer tools " + "are enabled. If None, auto-detected based on scrcpy & ffmpeg availability." + ), + ) + enable_ledger: bool = Field( + default=True, + description="Whether to enable video action ledger state tracking.", + ) + model_config = {"extra": "allow"} + + +class PlannerValidationConfig(BaseModel): + """Configuration specific to Planner validation of milestone changes.""" + + enabled: bool = Field( + default=False, + description="Whether to enable async planner validation when task plan milestones change.", + ) + similarity_threshold: float = Field( + default=0.85, + description="Difflib similarity threshold below which milestone text changes trigger re-validation.", + ) + model_config = {"extra": "allow"} + + +class CommitteeConfig(BaseModel): + """Configuration specific to Multi-Agent Committee / Council.""" + + enabled: bool = Field( + default=False, + description="Whether to equip Operator with the ask_committee tool.", + ) + debate_rounds: int = Field( + default=2, + ge=1, + le=5, + description="Number of debate rounds among council members (1-5).", + ) + model_config = {"extra": "allow"} + + +class CheckerConfig(BaseModel): + """Configuration specific to Checker subgoal verification and rollback.""" + + enabled: bool = Field( + default=False, + description="Whether to enable visual verification and snapshot rollback by the Checker agent.", + ) + max_iterations: int = Field( + default=20, + ge=1, + le=50, + description="Maximum verification iterations for Checker LLM reasoning and inspection.", + ) + max_chat_rounds: int = Field( + default=4, + ge=1, + le=10, + description="Maximum debate/rebuttal rounds between Operator and Checker.", + ) + model_config = {"extra": "allow"} + + +class OutputterConfig(BaseModel): + """Configuration specific to Outputter post-execution report synthesis agent.""" + + enabled: bool = Field( + default=True, + description="Whether Outputter agent is mounted for post-execution report/structured output synthesis.", + ) + force_synthesis: bool = Field( + default=False, + description="Whether to force Outputter synthesis even when no structured schema is explicitly requested.", + ) + model_config = {"extra": "allow"} + + +class ExplorerConfig(BaseModel): + """Configuration specific to UI Explorer visual sub-agent (Flash / Pro / Ultra).""" + + default_version: ExplorerVersion = Field( + default=DEFAULT_EXPLORER_VERSION, + description="Default Explorer version mode ('flash', 'pro', or 'ultra').", + ) + flash_mode: ExplorerVersion = Field( + default="flash", + description="Explorer version mode when main system runs under Flash profile (FlashRunner).", + ) + pro_mode: ExplorerVersion = Field( + default="flash", + description="Explorer version mode when main system runs under Pro profile (Graph / Operator / Validator).", + ) + caching: bool = Field( + default=True, + description="Whether to enable context caching for multi-turn pro/ultra Explorer.", + ) + model_config = {"extra": "allow"} + + def resolve( + self, + explicit_version: str | None = None, + agent_name: str | None = None, + profile: str | None = None, + per_agent_overrides: dict[str, str] | None = None, + ) -> ExplorerVersion: + """Resolves active Explorer version with clean precedence rules.""" + if explicit_version: + v = str(explicit_version).strip().lower() + if v in ("flash", "pro", "ultra"): + return v # type: ignore + + env_v = os.getenv("ARTEMIS_EXPLORER_VERSION", "").strip().lower() + if env_v in ("flash", "pro", "ultra"): + return env_v # type: ignore + + if per_agent_overrides and agent_name and agent_name in per_agent_overrides: + v = str(per_agent_overrides[agent_name]).strip().lower() + if v in ("flash", "pro", "ultra"): + return v # type: ignore + + if profile in ("flash", "flash_runner") or agent_name in ("flash", "flash_runner"): + return self.flash_mode + if profile in ("pro", "operator", "validator") or agent_name in ( + "pro", + "operator", + "validator", + ): + return self.pro_mode + + return self.default_version + + +class StepSummarizerConfig(BaseModel): + """Configuration specific to Flash asynchronous step visual summarization.""" + + enabled: bool = Field( + default=True, + description="Whether to asynchronously summarize historical steps to replace pruned images.", + ) + model: str = Field( + default="gemini-2.5-flash-lite", + description="Lightweight model used for background step state summarization.", + ) + prune_history_xml: bool = Field( + default=True, + description="Whether to prune outdated heavy UI XML trees from historical turns.", + ) + model_config = {"extra": "allow"} + + +class FlashProfileConfig(BaseModel): + """Configuration options specific to the โšก Flash execution profile (FlashRunner/ReactiveRunner).""" + + max_turns: int = Field( + default=30, + ge=1, + le=100, + description="Maximum reactive turns for Flash execution before timeout.", + ) + explorer_mode: ExplorerVersion = Field( + default="flash", + description="Perception mode for Flash runner ('flash' for 1-shot detection).", + ) + step_summarizer: StepSummarizerConfig = Field( + default_factory=StepSummarizerConfig, + description="Asynchronous visual context compressor settings.", + ) + model_config = {"extra": "allow"} + + +class ProExplorerConfig(BaseModel): + """Explorer perception settings under Pro profile.""" + + mode: ExplorerVersion = Field( + default="flash", + description="Explorer perception mode under Pro profile ('flash', 'pro', or 'ultra').", + ) + caching: bool = Field( + default=True, + description="Whether to enable context caching for multi-turn Pro/Ultra exploration.", + ) + model_config = {"extra": "allow"} + + +class ProProfileConfig(BaseModel): + """Configuration options specific to the ๐Ÿš€ Pro execution profile (LangGraph / Multi-Agent Closed-Loop).""" + + explorer: ProExplorerConfig = Field( + default_factory=ProExplorerConfig, + description="Explorer perception settings for Pro profile.", + ) + planner_validation: PlannerValidationConfig = Field( + default_factory=PlannerValidationConfig, + description="Planner milestone consistency validation settings.", + ) + committee: CommitteeConfig = Field( + default_factory=CommitteeConfig, + description="Multi-Agent Committee council debate settings.", + ) + checker: CheckerConfig = Field( + default_factory=CheckerConfig, + description="Checker visual verification and rollback settings.", + ) + video_analyzer: VideoAnalyzerConfig = Field( + default_factory=VideoAnalyzerConfig, + description="Screen video action ledger tracking settings.", + ) + model_config = {"extra": "allow"} + + +class AgentGlobalConfig(BaseModel): + """Global configuration parsed from artemis.jsonc / agent_config.json.""" + + # โšก Profile-categorized configurations + flash: FlashProfileConfig = Field( + default_factory=FlashProfileConfig, + description="Options for โšก Flash reactive runner execution.", + ) + pro: ProProfileConfig = Field( + default_factory=ProProfileConfig, + description="Options for ๐Ÿš€ Pro multi-agent graph execution.", + ) + + # ๐ŸŒ Core Component Configurations + explorer: ExplorerConfig = Field( + default_factory=ExplorerConfig, + description="UI Explorer sub-agent runtime options.", + ) + explorer_versions: dict[str, str] = Field( + default_factory=lambda: {"operator": "flash", "validator": "flash"}, + description="Per-agent mapping of UI explorer version (e.g. flash, pro, ultra).", + ) + planner_validation: PlannerValidationConfig = Field( + default_factory=PlannerValidationConfig, + description="Planner milestone validation runtime options.", + ) + committee: CommitteeConfig = Field( + default_factory=CommitteeConfig, + description="Committee council runtime options.", + ) + checker: CheckerConfig = Field( + default_factory=CheckerConfig, + description="Checker subgoal verification and rollback runtime options.", + ) + outputter: OutputterConfig = Field( + default_factory=OutputterConfig, + description="Outputter synthesis agent runtime options.", + ) + blacklisted_tools: dict[str, list[str]] = Field( + default_factory=dict, + description="List of blacklisted tools per agent.", + ) + video_analyzer: VideoAnalyzerConfig = Field( + default_factory=VideoAnalyzerConfig, + description="Video analyzer runtime options.", + ) + + model_config = {"extra": "allow"} + + @model_validator(mode="before") + @classmethod + def synchronize_profiles_and_components(cls, data: Any) -> Any: + """Bidirectionally synchronize between categorized 'flash'/'pro' blocks and component fields.""" + if not isinstance(data, dict): + return data + + # 1. Sync from pro block to root components if provided + if "pro" in data and isinstance(data["pro"], dict): + pro_data = data["pro"] + if "checker" in pro_data and "checker" not in data: + data["checker"] = pro_data["checker"] + if "committee" in pro_data and "committee" not in data: + data["committee"] = pro_data["committee"] + if "planner_validation" in pro_data and "planner_validation" not in data: + data["planner_validation"] = pro_data["planner_validation"] + if "video_analyzer" in pro_data and "video_analyzer" not in data: + data["video_analyzer"] = pro_data["video_analyzer"] + if "explorer" in pro_data and isinstance(pro_data["explorer"], dict): + exp_pro = pro_data["explorer"] + if "explorer" not in data: + data["explorer"] = {} + if isinstance(data["explorer"], dict): + if "mode" in exp_pro and "pro_mode" not in data["explorer"]: + data["explorer"]["pro_mode"] = exp_pro["mode"] + if "caching" in exp_pro and "caching" not in data["explorer"]: + data["explorer"]["caching"] = exp_pro["caching"] + + # 2. Sync from flash block to explorer.flash_mode + if "flash" in data and isinstance(data["flash"], dict): + flash_data = data["flash"] + if "explorer_mode" in flash_data: + if "explorer" not in data: + data["explorer"] = {} + if isinstance(data["explorer"], dict) and "flash_mode" not in data["explorer"]: + data["explorer"]["flash_mode"] = flash_data["explorer_mode"] + + # 3. Sync from root components to pro/flash blocks if pro/flash block is omitted + if "checker" in data and ("pro" not in data or "checker" not in data.get("pro", {})): + data.setdefault("pro", {})["checker"] = data["checker"] + if "committee" in data and ("pro" not in data or "committee" not in data.get("pro", {})): + data.setdefault("pro", {})["committee"] = data["committee"] + if "planner_validation" in data and ( + "pro" not in data or "planner_validation" not in data.get("pro", {}) + ): + data.setdefault("pro", {})["planner_validation"] = data["planner_validation"] + if "video_analyzer" in data and ( + "pro" not in data or "video_analyzer" not in data.get("pro", {}) + ): + data.setdefault("pro", {})["video_analyzer"] = data["video_analyzer"] + + # 4. Environment variable overrides (highest precedence for runtime switches) + def _parse_bool(var_name: str) -> bool | None: + v = os.getenv(var_name) + if v is None: + return None + v_clean = v.strip().lower() + if v_clean in ("1", "true", "yes", "on"): + return True + if v_clean in ("0", "false", "no", "off"): + return False + return None + + chk_env = _parse_bool("ARTEMIS_CHECKER_ENABLED") + if chk_env is not None: + data.setdefault("checker", {})["enabled"] = chk_env + data.setdefault("pro", {}).setdefault("checker", {})["enabled"] = chk_env + + comm_env = _parse_bool("ARTEMIS_COMMITTEE_ENABLED") + if comm_env is not None: + data.setdefault("committee", {})["enabled"] = comm_env + data.setdefault("pro", {}).setdefault("committee", {})["enabled"] = comm_env + + pv_env = _parse_bool("ARTEMIS_PLANNER_VALIDATION_ENABLED") + if pv_env is not None: + data.setdefault("planner_validation", {})["enabled"] = pv_env + data.setdefault("pro", {}).setdefault("planner_validation", {})["enabled"] = pv_env + + out_env = _parse_bool("ARTEMIS_OUTPUTTER_ENABLED") + if out_env is not None: + data.setdefault("outputter", {})["enabled"] = out_env + + vid_rec_env = _parse_bool("ARTEMIS_VIDEO_RECORDING_ENABLED") or _parse_bool( + "ARTEMIS_WITH_VIDEO_RECORDING_TOOLS" + ) + if vid_rec_env is not None: + data.setdefault("video_analyzer", {})["enabled"] = vid_rec_env + data.setdefault("pro", {}).setdefault("video_analyzer", {})["enabled"] = vid_rec_env + + vid_env = _parse_bool("ARTEMIS_VIDEO_LEDGER_ENABLED") + if vid_env is not None: + data.setdefault("video_analyzer", {})["enable_ledger"] = vid_env + data.setdefault("pro", {}).setdefault("video_analyzer", {})["enable_ledger"] = vid_env + + return data + + def get_explorer_version( + self, + explicit_version: str | None = None, + agent_name: str | None = "operator", + profile: str | None = None, + ) -> ExplorerVersion: + """Resolve Explorer version from this configuration.""" + return self.explorer.resolve( + explicit_version=explicit_version, + agent_name=agent_name, + profile=profile, + per_agent_overrides=self.explorer_versions, + ) + + +def resolve_explorer_version( + ctx: Any | None = None, + explicit_version: str | None = None, + agent_or_profile_name: str | None = "operator", +) -> ExplorerVersion: + """Resolves active Explorer version based on explicit args, env vars, agent_config, and execution setup.""" + if explicit_version: + v = str(explicit_version).strip().lower() + if v in ("flash", "pro", "ultra"): + return v # type: ignore + + env_v = os.getenv("ARTEMIS_EXPLORER_VERSION", "").strip().lower() + if env_v in ("flash", "pro", "ultra"): + return env_v # type: ignore + + if ctx: + agent_cfg = getattr(ctx, "agent_config", None) + exec_setup = getattr(ctx, "execution_setup", None) + + if agent_cfg: + if hasattr(agent_cfg, "get_explorer_version"): + return agent_cfg.get_explorer_version( + explicit_version=explicit_version, + agent_name=agent_or_profile_name, + ) + if hasattr(agent_cfg, "explorer") and hasattr(agent_cfg.explorer, "resolve"): + overrides = getattr(agent_cfg, "explorer_versions", None) + return agent_cfg.explorer.resolve( + explicit_version=explicit_version, + agent_name=agent_or_profile_name, + per_agent_overrides=overrides, + ) + + if ( + exec_setup + and hasattr(exec_setup, "explorer") + and hasattr(exec_setup.explorer, "resolve") + ): + overrides = getattr(exec_setup, "explorer_versions", None) + return exec_setup.explorer.resolve( + explicit_version=explicit_version, + agent_name=agent_or_profile_name, + per_agent_overrides=overrides, + ) + + return "flash" + + +def load_agent_config( + config_path: Path | None = None, +) -> AgentGlobalConfig: + """Load and validate AgentGlobalConfig from artemis.jsonc or agent_config.json. + + Args: + config_path: Optional explicit file path to configuration file. + + Returns: + Validated AgentGlobalConfig instance. + """ + resolved_path = config_path + if not resolved_path: + for candidate in ("artemis.jsonc", "artemis.json", AGENT_CONFIG_FILENAME): + try: + resolved_path = get_config_path(candidate) + break + except FileNotFoundError: + continue + + if not resolved_path: + resolved_path = get_config_path(AGENT_CONFIG_FILENAME, ROOT_DIR / AGENT_CONFIG_FILENAME) + + try: + with open(resolved_path, encoding="utf-8") as f: + data = load_jsonc(f) + if isinstance(data, dict) and "agent" in data: + return AgentGlobalConfig.model_validate(data["agent"]) + return AgentGlobalConfig.model_validate(data) + except Exception as e: + logger.error(f"Failed to load agent config from '{resolved_path}': {e}") + raise ValueError(f"Failed to parse '{resolved_path}': {e}") from e diff --git a/artemis/config/constants.py b/artemis/config/constants.py new file mode 100644 index 0000000..2fdb3a3 --- /dev/null +++ b/artemis/config/constants.py @@ -0,0 +1,142 @@ +# 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. + +"""Centralized constants, environment variable names, and type definitions for ARTEMIS.""" + +from typing import Literal + +# ============================================================================== +# Environment Variable Keys +# ============================================================================== + +# Application & Workspace Paths +ENV_ARTEMIS_APP_DIR = "ARTEMIS_APP_DIR" +ENV_ANTIGRAVITY_APP_DIR = "ANTIGRAVITY_APP_DIR" +ENV_ARTEMIS_USE_USER_DIR = "ARTEMIS_USE_USER_DIR" +ENV_ARTEMIS_TRACES_DIR = "ARTEMIS_TRACES_DIR" +ENV_DATA_ENGINE_DB_PATH = "DATA_ENGINE_DB_PATH" + +# Cloud Brain & Distributed Execution +ENV_ARTEMIS_CLOUD_MODE = "ARTEMIS_CLOUD_MODE" +ENV_ARTEMIS_CLOUD_SESSION_ID = "ARTEMIS_CLOUD_SESSION_ID" +ENV_ARTEMIS_CLOUD_GATEWAY_URL = "ARTEMIS_CLOUD_GATEWAY_URL" +ENV_ARTEMIS_TENANT_ID = "ARTEMIS_TENANT_ID" +ENV_ARTEMIS_TENANT_TOKEN = "ARTEMIS_TENANT_TOKEN" +ENV_ARTEMIS_EDGE_PORT = "ARTEMIS_EDGE_PORT" + +# Runtime State & IPC +ENV_ARTEMIS_IPC_PORT = "ARTEMIS_IPC_PORT" +ENV_ANTIGRAVITY_LS_ADDRESS = "ANTIGRAVITY_LS_ADDRESS" +ENV_ARTEMIS_MCP_SERVER = "ARTEMIS_MCP_SERVER" + +# Device & ADB +ENV_ADB_DEVICE_SERIAL = "ADB_DEVICE_SERIAL" +ENV_ARTEMIS_DEVICE_ID = "ARTEMIS_DEVICE_ID" +ENV_ADB_HOST = "ADB_HOST" +ENV_ADB_PORT = "ADB_PORT" +ENV_ADB_SERVER_SOCKET = "ADB_SERVER_SOCKET" + +# LLM & Vision OCR API Keys +ENV_GOOGLE_API_KEY = "GOOGLE_API_KEY" +ENV_GEMINI_API_KEY = "GEMINI_API_KEY" +ENV_GCP_API_KEY = "GCP_API_KEY" +ENV_OCR_API_KEY = "OCR_API_KEY" +ENV_VISION_API_KEY = "VISION_API_KEY" +ENV_API_KEY = "API_KEY" +ENV_OPENAI_API_KEY = "OPENAI_API_KEY" +ENV_OPENAI_BASE_URL = "OPENAI_BASE_URL" +ENV_ANTHROPIC_API_KEY = "ANTHROPIC_API_KEY" +ENV_OPEN_ROUTER_API_KEY = "OPEN_ROUTER_API_KEY" +ENV_XAI_API_KEY = "XAI_API_KEY" + +# Explorer & Agent Defaults +ENV_ARTEMIS_EXPLORER_VERSION = "ARTEMIS_EXPLORER_VERSION" +ENV_ARTEMIS_DEFAULT_PROFILE = "ARTEMIS_DEFAULT_PROFILE" +ENV_ARTEMIS_DEFAULT_MODEL = "ARTEMIS_DEFAULT_MODEL" +ENV_ARTEMIS_USE_FILE_API = "ARTEMIS_USE_FILE_API" + +# Debugging & Output Paths +ENV_KEEP_VIDEOS = "KEEP_VIDEOS" +ENV_ARTEMIS_DEBUG = "ARTEMIS_DEBUG" +ENV_EVENTS_OUTPUT_PATH = "EVENTS_OUTPUT_PATH" +ENV_RESULTS_OUTPUT_PATH = "RESULTS_OUTPUT_PATH" + + +# ============================================================================== +# Configuration Filenames +# ============================================================================== + +ARTEMIS_CONFIG_FILENAME = "artemis.jsonc" +LLM_CONFIG_FILENAME = "llm-config.json" +LLM_CONFIG_OVERRIDE_FILENAME = "llm-config.override.jsonc" +AGENT_CONFIG_FILENAME = "agent_config.json" +DATA_ENGINE_DB_FILENAME = "data_engine.db" +IPC_PORT_FILENAME = ".artemis_ipc_port" +LS_ADDRESS_FILENAME = ".jetski_ls_address" +DOTENV_FILENAME = ".env" + + +# ============================================================================== +# Default Values & Network Ports +# ============================================================================== + +DEFAULT_ADB_HOST = "127.0.0.1" +DEFAULT_ADB_PORT = 5037 +DEFAULT_GATEWAY_PORT = 8000 +DEFAULT_EDGE_MONITOR_PORT = 8000 +DEFAULT_ACTION_TIMEOUT = 30.0 +DEFAULT_STREAM_PING_INTERVAL = 15.0 + +DEFAULT_PROFILE = "pro" +DEFAULT_MODEL = "gemini-2.5-flash" +DEFAULT_EXPLORER_VERSION: Literal["flash", "pro", "ultra"] = "flash" + + +# ============================================================================== +# Literals and Type Aliases +# ============================================================================== + +LLMProvider = Literal[ + "openai", "google", "openrouter", "xai", "vertexai", "anthropic", "ollama", "vllm", "custom" +] +ExplorerVersion = Literal["flash", "pro", "ultra"] + +LLMUtilsNode = Literal[ + "outputter", + "hopper", + "video_analyzer", + "object_detector", +] +LLMUtilsNodeWithFallback = LLMUtilsNode + +AgentNode = Literal[ + "planner", + "summarizer", + "operator", + "operator_summarizer", + "log_reader_sub_agent", + "log_analyzer", + "diagnoser", + "validator_failure_analyzer", + "checker", + "planner_avatar", + "history_analyzer_expert", + "diagnoser_expert", + "explorer", + "history_analyzer", + "validator_pixel_safety_net", + "validator", + "output_analyzer", +] +AgentNodeWithFallback = AgentNode diff --git a/artemis/config/llm.py b/artemis/config/llm.py new file mode 100644 index 0000000..de2177d --- /dev/null +++ b/artemis/config/llm.py @@ -0,0 +1,364 @@ +# 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. + +"""LLM provider, model hierarchy, fallback chaining, and configuration loaders.""" + +import os +from pathlib import Path +from typing import Any, Literal + +import google.auth +from google.auth.exceptions import DefaultCredentialsError +from pydantic import BaseModel, ValidationError + +from artemis.config.constants import ( + LLM_CONFIG_FILENAME, + AgentNode, + LLMProvider, + LLMUtilsNode, +) +from artemis.config.paths import ROOT_DIR, get_config_path +from artemis.config.settings import settings +from artemis.utils.file import load_jsonc +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +def validate_vertex_ai_credentials() -> None: + """Validate Google Application Default Credentials for VertexAI provider.""" + try: + _, project = google.auth.default() + if not project: + raise Exception("VertexAI requires a Google Cloud project to be set.") + except DefaultCredentialsError as e: + raise Exception( + f"VertexAI requires valid Google Application Default Credentials (ADC): {e}" + ) + + +class _CyFunctionDetectorMeta(type): + def __instancecheck__(self, instance): + name = type(instance).__name__ + return ( + name + in ( + "cyfunction", + "cython_function_or_method", + "builtin_function_or_method", + ) + or "cyfunction" in name.lower() + ) + + +class CyFunctionDetector(metaclass=_CyFunctionDetectorMeta): + """Utility type detector for Cython and C-extension functions in Pydantic.""" + + pass + + +class LLM(BaseModel): + """Base model representing an LLM model provider and runtime parameters.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + provider: LLMProvider + model: str + temperature: float | None = None + thinking_budget: int | None = None + thinking_level: Literal["minimal", "low", "medium", "high"] | None = None + reasoning_effort: Literal["none", "low", "medium", "high"] | None = None + include_thoughts: bool | None = None + enable_grounding: bool | None = None + + def validate_provider(self, name: str) -> None: + """Ensure the required API key or credentials exist in settings for this provider.""" + if self.provider == "openai": + if not settings.OPENAI_API_KEY: + raise Exception(f"{name} requires OPENAI_API_KEY in .env") + elif self.provider == "google": + if not settings.GOOGLE_API_KEY: + raise Exception(f"{name} requires GOOGLE_API_KEY in .env") + elif self.provider == "vertexai": + validate_vertex_ai_credentials() + elif self.provider == "anthropic": + if not (settings.ANTHROPIC_API_KEY or os.environ.get("ANTHROPIC_API_KEY")): + raise Exception(f"{name} requires ANTHROPIC_API_KEY in .env") + elif self.provider == "openrouter": + if not settings.OPEN_ROUTER_API_KEY: + raise Exception(f"{name} requires OPEN_ROUTER_API_KEY in .env") + elif self.provider == "xai": + if not settings.XAI_API_KEY: + raise Exception(f"{name} requires XAI_API_KEY in .env") + + def __str__(self) -> str: + return f"{self.provider}/{self.model}" + + +class LLMWithFallback(LLM): + """LLM configuration with automatic secondary fallback and timeout specs.""" + + fallback: LLM + fix_model: str | None = None + timeout: float | None = None + + def __str__(self) -> str: + return f"{self.provider}/{self.model} (fallback: {self.fallback})" + + +class LLMConfigUtils(BaseModel): + """Configuration container for auxiliary utility agents/nodes.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + outputter: LLMWithFallback + hopper: LLMWithFallback + video_analyzer: LLMWithFallback | None = None + object_detector: LLMWithFallback | None = None + + +class LLMConfig(BaseModel): + """Comprehensive LLM configuration mapping every node to primary/fallback models.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + planner: LLMWithFallback + utils: LLMConfigUtils + summarizer: LLMWithFallback + operator: LLMWithFallback + operator_summarizer: LLMWithFallback + log_reader_sub_agent: LLMWithFallback + log_analyzer: LLMWithFallback + diagnoser: LLMWithFallback + validator_failure_analyzer: LLMWithFallback + checker: LLMWithFallback + planner_avatar: LLMWithFallback + history_analyzer_expert: LLMWithFallback + diagnoser_expert: LLMWithFallback + explorer: LLMWithFallback + history_analyzer: LLMWithFallback | None = None + validator_pixel_safety_net: LLMWithFallback | None = None + output_analyzer: LLMWithFallback | None = None + + def validate_providers(self) -> None: + """Validate credentials across all configured agent nodes.""" + self.planner.validate_provider("Planner") + self.utils.outputter.validate_provider("Outputter") + self.utils.hopper.validate_provider("Hopper") + if self.utils.video_analyzer: + self.utils.video_analyzer.validate_provider("VideoAnalyzer") + if self.utils.object_detector: + self.utils.object_detector.validate_provider("ObjectDetector") + self.summarizer.validate_provider("Summarizer") + self.operator.validate_provider("Operator") + self.operator_summarizer.validate_provider("OperatorSummarizer") + self.log_reader_sub_agent.validate_provider("LogReaderSubAgent") + self.log_analyzer.validate_provider("LogAnalyzer") + self.diagnoser.validate_provider("Diagnoser") + self.validator_failure_analyzer.validate_provider("ValidatorFailureAnalyzer") + self.checker.validate_provider("Checker") + self.planner_avatar.validate_provider("PlannerAvatar") + self.history_analyzer_expert.validate_provider("HistoryAnalyzerExpert") + self.diagnoser_expert.validate_provider("DiagnoserExpert") + self.explorer.validate_provider("Explorer") + if self.history_analyzer: + self.history_analyzer.validate_provider("HistoryAnalyzer") + if self.validator_pixel_safety_net: + self.validator_pixel_safety_net.validate_provider("ValidatorPixelSafetyNet") + if self.output_analyzer: + self.output_analyzer.validate_provider("OutputAnalyzer") + + def __str__(self) -> str: + return f""" +๐Ÿ“ƒ Planner: {self.planner} +๐Ÿงฉ Utils: + ๐Ÿ”ฝ Hopper: {self.utils.hopper} + ๐Ÿ“ Outputter: {self.utils.outputter} + ๐ŸŽฌ Video Analyzer: {self.utils.video_analyzer or "Not configured"} + ๐Ÿ‘๏ธ Object Detector: {self.utils.object_detector or "Not configured"} +""" + + def get_agent(self, item: AgentNode) -> LLMWithFallback: + """Retrieve model configuration for a specific agent node with sensible defaults.""" + val = getattr(self, item) + if val is None: + if item == "history_analyzer": + return self.operator + elif item == "validator_pixel_safety_net": + return LLMWithFallback( + provider="google", + model="gemini-3.5-flash-lite", + temperature=0.0, + fallback=LLM( + provider="google", + model="gemini-3.1-flash-lite", + temperature=0.0, + ), + ) + elif item == "output_analyzer": + return self.log_analyzer + return val + + def get_utils(self, item: LLMUtilsNode) -> LLMWithFallback: + """Retrieve model configuration for a specific utility node.""" + value = getattr(self.utils, item) + if value is None: + raise ValueError( + f"Utils '{item}' is not configured. Please add it to your LLM " + "config or enable it via AgentConfigBuilder." + ) + return value + + +def _expand_default_into_nodes(config_dict: dict) -> dict: + """Expand unified config format with 'default' and 'nodes' into full LLMConfig schema.""" + if "planner" in config_dict and "utils" in config_dict: + return config_dict + + default_model_cfg = config_dict.get( + "default", + { + "provider": "google", + "model": "gemini-3.6-flash", + "fallback": { + "provider": "google", + "model": "gemini-3.5-flash", + }, + }, + ) + + nodes_override = config_dict.get("nodes", {}) + + all_agent_nodes = [ + "planner", + "summarizer", + "operator", + "operator_summarizer", + "log_reader_sub_agent", + "log_analyzer", + "diagnoser", + "validator_failure_analyzer", + "checker", + "planner_avatar", + "history_analyzer_expert", + "diagnoser_expert", + "explorer", + ] + + all_utils_nodes = [ + "outputter", + "hopper", + "video_analyzer", + "object_detector", + ] + + result: dict[str, Any] = {} + for node in all_agent_nodes: + node_cfg = dict(default_model_cfg) + if node in nodes_override: + for k, v in nodes_override[node].items(): + if isinstance(v, dict) and isinstance(node_cfg.get(k), dict): + node_cfg[k] = {**node_cfg[k], **v} + else: + node_cfg[k] = v + result[node] = node_cfg + + utils_dict: dict[str, Any] = {} + for util in all_utils_nodes: + util_cfg = dict(default_model_cfg) + if util in nodes_override: + for k, v in nodes_override[util].items(): + if isinstance(v, dict) and isinstance(util_cfg.get(k), dict): + util_cfg[k] = {**util_cfg[k], **v} + else: + util_cfg[k] = v + utils_dict[util] = util_cfg + result["utils"] = utils_dict + + return result + + +def parse_llm_config() -> LLMConfig: + """Parse and instantiate LLMConfig from artemis.jsonc or llm-config.json.""" + config_path = None + for candidate in ("artemis.jsonc", "artemis.json", LLM_CONFIG_FILENAME): + try: + config_path = get_config_path(candidate) + break + except FileNotFoundError: + continue + + if not config_path: + config_path = get_config_path(LLM_CONFIG_FILENAME, ROOT_DIR / LLM_CONFIG_FILENAME) + + try: + with open(config_path, encoding="utf-8") as f: + config_dict = load_jsonc(f) + expanded_dict = _expand_default_into_nodes(config_dict) + return LLMConfig.model_validate(expanded_dict) + except Exception as e: + logger.error(f"Failed to load or parse llm config: {config_path}. Error: {e}") + raise + + +def initialize_llm_config() -> LLMConfig: + """Parse and validate credentials for LLMConfig.""" + llm_config = parse_llm_config() + llm_config.validate_providers() + logger.success("LLM config initialized") + return llm_config + + +def get_default_llm_config() -> LLMConfig: + """Returns default LLMConfig parsed from standard configuration file.""" + return parse_llm_config() + + +def deep_merge_llm_config(base: LLMConfig, overrides: dict) -> LLMConfig: + """Recursively merge dictionary overrides into an existing LLMConfig object.""" + base_dict = base.model_dump() + + def merge(d1: dict, d2: dict) -> None: + for k, v in d2.items(): + if k in d1 and isinstance(d1[k], dict) and isinstance(v, dict): + merge(d1[k], v) + else: + d1[k] = v + + merge(base_dict, overrides) + return LLMConfig.model_validate(base_dict) + + +def load_llm_config_override(path: Path | str) -> LLMConfig: + """Load custom LLM configuration JSON/JSONC overrides on top of default configuration.""" + default_config = get_default_llm_config() + + resolved_path = Path(path) + if not resolved_path.exists(): + try: + resolved_path = get_config_path(str(path)) + except Exception: + pass + + override_config_dict = {} + if resolved_path.exists(): + logger.info(f"Loading custom LLM config from {resolved_path.resolve()}...") + with open(resolved_path, encoding="utf-8") as f: + override_config_dict = load_jsonc(f) + else: + logger.warning(f"Custom LLM config not found at {path} - using default config") + + try: + return deep_merge_llm_config(default_config, override_config_dict) + except ValidationError as e: + logger.error(f"Invalid LLM config: {e}") + logger.info("Falling back to default config") + return default_config diff --git a/artemis/config/output.py b/artemis/config/output.py new file mode 100644 index 0000000..1955b9a --- /dev/null +++ b/artemis/config/output.py @@ -0,0 +1,159 @@ +# 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. + +"""Structured output schemas, event recording, and output file management.""" + +import json +import os +from pathlib import Path +from typing import Annotated, Any +import warnings + +from pydantic import BaseModel, Field, model_validator + +from artemis.config.constants import ( + ENV_EVENTS_OUTPUT_PATH, + ENV_RESULTS_OUTPUT_PATH, +) +from artemis.config.llm import CyFunctionDetector +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +class OutputConfig(BaseModel): + """Configuration for structured output formatting and schemas.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + structured_output: Annotated[ + type[BaseModel] | dict | None, + Field( + default=None, + description=( + "Optional structured schema (as a BaseModel or dict) to shape the output. " + "If provided, it takes precedence over 'output_description'." + ), + ), + ] + output_description: Annotated[ + str | None, + Field( + default=None, + description=( + "Optional natural language description of the expected output format. " + "Used only if 'structured_output' is not provided. " + "Example: 'Output a JSON with 3 keys: color, price, websiteUrl'." + ), + ), + ] + enable_outputter: Annotated[ + bool | None, + Field( + default=None, + description="Explicitly enable or disable Outputter post-execution synthesis.", + ), + ] = None + force_synthesis: Annotated[ + bool, + Field( + default=False, + description="Whether to run Outputter synthesis even without explicit output_description or schema.", + ), + ] = False + + def __str__(self) -> str: + s_builder = "" + if self.structured_output: + s_builder += f"Structured Output: {self.structured_output}\n" + if self.output_description: + s_builder += f"Output Description: {self.output_description}\n" + if self.output_description and self.structured_output: + s_builder += ( + "Both 'structured_output' and 'output_description' are provided. " + "'structured_output' will take precedence.\n" + ) + return s_builder + + @model_validator(mode="after") + def warn_if_both_outputs_provided(self) -> "OutputConfig": + if self.structured_output and self.output_description: + warnings.warn( + "Both 'structured_output' and 'output_description' are provided. " + "'structured_output' will take precedence.", + stacklevel=2, + ) + return self + + def needs_structured_format(self, default_enabled: bool = True) -> bool: + """Returns True if Outputter should run to synthesize/structure the output.""" + if self.enable_outputter is False: + return False + if self.enable_outputter is True or self.force_synthesis: + return True + if not default_enabled: + return False + return bool(self.structured_output or self.output_description) + + +def prepare_output_files() -> tuple[str | None, str | None]: + """Validate and prepare events and results output files from environment variables.""" + events_output_path = os.getenv(ENV_EVENTS_OUTPUT_PATH) + results_output_path = os.getenv(ENV_RESULTS_OUTPUT_PATH) + + def validate_and_prepare_file(file_path: str) -> str | None: + if not file_path: + return None + + path_obj = Path(file_path) + + if path_obj.exists() and path_obj.is_dir(): + logger.error(f"Error: Path '{file_path}' points to an existing directory, not a file.") + return None + + if not path_obj.suffix or file_path.endswith(("/", "\\")): + logger.error(f"Error: Path '{file_path}' appears to be a directory path, not a file.") + return None + + try: + path_obj.parent.mkdir(parents=True, exist_ok=True) + path_obj.touch(exist_ok=True) + return file_path + except OSError as e: + logger.error(f"Error creating file '{file_path}': {e}") + return None + + validated_events_path = ( + validate_and_prepare_file(events_output_path) if events_output_path else None + ) + validated_results_path = ( + validate_and_prepare_file(results_output_path) if results_output_path else None + ) + + return validated_events_path, validated_results_path + + +def record_events(output_path: Path | None, events: list[str] | BaseModel | Any) -> None: + """Serialize and write execution events to disk.""" + if not output_path: + return + + if isinstance(events, str): + events_content = events + elif isinstance(events, BaseModel): + events_content = events.model_dump_json(indent=2) + else: + events_content = json.dumps(events, indent=2) + + with open(output_path, "w", encoding="utf-8") as f: + f.write(events_content) diff --git a/artemis/config/paths.py b/artemis/config/paths.py new file mode 100644 index 0000000..0015d82 --- /dev/null +++ b/artemis/config/paths.py @@ -0,0 +1,170 @@ +# 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. + +"""Centralized workspace, configuration, temporary files, and directory management.""" + +import os +from pathlib import Path +import sys + +from artemis.config.constants import ( + DATA_ENGINE_DB_FILENAME, + ENV_ANTIGRAVITY_APP_DIR, + ENV_ARTEMIS_APP_DIR, + ENV_ARTEMIS_TRACES_DIR, + ENV_ARTEMIS_USE_USER_DIR, + IPC_PORT_FILENAME, + LS_ADDRESS_FILENAME, +) + +# Project root directory (3 levels up from artemis/config/paths.py) +ROOT_DIR = Path(__file__).resolve().parent.parent.parent +CONFIG_DIR = ROOT_DIR / "config" + + +def is_frozen_bundle() -> bool: + """Check if running inside a PyInstaller / Nuitka standalone compiled binary.""" + return getattr(sys, "frozen", False) and hasattr(sys, "_MEIPASS") + + +def get_app_dir() -> Path: + """Returns the central application directory for ARTEMIS configurations, DBs, and traces. + + Order of precedence: + 1. ARTEMIS_APP_DIR or ANTIGRAVITY_APP_DIR environment variable. + 2. Legacy ~/.gemini/jetski directory if it exists. + 3. Default ~/.artemis directory. + """ + env_dir = os.getenv(ENV_ARTEMIS_APP_DIR) or os.getenv(ENV_ANTIGRAVITY_APP_DIR) + if env_dir: + path = Path(env_dir) + path.mkdir(parents=True, exist_ok=True) + return path + + legacy_dir = Path(os.path.expanduser("~/.gemini/jetski")) + if legacy_dir.exists(): + return legacy_dir + + default_dir = Path(os.path.expanduser("~/.artemis")) + default_dir.mkdir(parents=True, exist_ok=True) + return default_dir + + +def get_default_traces_path() -> Path: + """Returns default traces directory. + + Uses user app dir when running as a frozen binary or when app dir env is set. + """ + env_traces = os.getenv(ENV_ARTEMIS_TRACES_DIR) + if env_traces: + traces_dir = Path(env_traces) + elif ( + is_frozen_bundle() + or os.getenv(ENV_ARTEMIS_USE_USER_DIR) == "true" + or os.getenv(ENV_ARTEMIS_APP_DIR) + or os.getenv(ENV_ANTIGRAVITY_APP_DIR) + ): + traces_dir = get_app_dir() / "traces" + else: + traces_dir = ROOT_DIR / "traces" + + traces_dir.mkdir(parents=True, exist_ok=True) + return traces_dir + + +def get_traces_dir() -> Path: + """Alias for get_default_traces_path.""" + return get_default_traces_path() + + +def get_temp_dir(subfolder: str | None = None) -> Path: + """Returns a unified temporary directory for ARTEMIS runtime operations (e.g., screenshots, recordings). + + Args: + subfolder: Optional subfolder name under the temp directory. + + Returns: + Path object pointing to the initialized temporary directory. + """ + base_temp = get_app_dir() / "temp" + target_dir = base_temp / subfolder if subfolder else base_temp + target_dir.mkdir(parents=True, exist_ok=True) + return target_dir + + +def get_data_engine_db_path() -> Path: + """Returns the unified SQLite database path for ARTEMIS DataEngine.""" + return get_default_traces_path() / DATA_ENGINE_DB_FILENAME + + +def get_ipc_port_file() -> Path: + """Returns the location of the IPC port synchronization file.""" + app_dir_file = get_app_dir() / IPC_PORT_FILENAME + if app_dir_file.exists(): + return app_dir_file + return ROOT_DIR / IPC_PORT_FILENAME + + +def get_ls_address_file() -> Path: + """Returns the location of the Language Server address synchronization file.""" + app_dir_file = get_app_dir() / LS_ADDRESS_FILENAME + if app_dir_file.exists(): + return app_dir_file + return ROOT_DIR / LS_ADDRESS_FILENAME + + +def get_config_path(filename: str, default_bundled_path: Path | None = None) -> Path: + """Resolve a configuration file path across env vars, config dir, project dir, app dir, and bundles. + + Resolution Order: + 1. Environment variable override: `ARTEMIS_` + 2. Central `config/` directory: `/config/` + 3. Project root directory: `/` + 4. Custom default bundled path (if provided and exists) + 5. User application directory: `/` + """ + env_var = f"ARTEMIS_{filename.upper().replace('.', '_').replace('-', '_')}" + env_path_str = os.getenv(env_var) + if env_path_str: + env_path = Path(env_path_str) + if env_path.exists(): + return env_path + + # Check config/ directory + config_dir_path = CONFIG_DIR / filename + if config_dir_path.exists(): + return config_dir_path + + # Check project root directory + project_path = ROOT_DIR / filename + if project_path.exists(): + return project_path + + # Check bundled fallback + if default_bundled_path and default_bundled_path.exists(): + return default_bundled_path + + # Check user app directory + app_dir_path = get_app_dir() / filename + if app_dir_path.exists(): + return app_dir_path + + raise FileNotFoundError( + f"Configuration file '{filename}' not found in config dir ({config_dir_path}), " + f"project root ({project_path}), or app dir ({app_dir_path})." + ) + + +GLOBAL_APP_DIR = get_app_dir() +GLOBAL_JETSKI_DIR = GLOBAL_APP_DIR # Backward-compatibility alias diff --git a/artemis/config/runtime.py b/artemis/config/runtime.py new file mode 100644 index 0000000..4f30867 --- /dev/null +++ b/artemis/config/runtime.py @@ -0,0 +1,142 @@ +# 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. + +"""Runtime state coordination, temporary file lifecycles, and IPC synchronization.""" + +import os +from pathlib import Path +import time + +from artemis.config.constants import ( + ENV_ANTIGRAVITY_LS_ADDRESS, + ENV_ARTEMIS_IPC_PORT, +) +from artemis.config.paths import ( + get_ipc_port_file, + get_ls_address_file, + get_temp_dir, +) +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +def read_ipc_port() -> int | None: + """Read the active ARTEMIS IPC port from environment or temporary state file.""" + env_port = os.getenv(ENV_ARTEMIS_IPC_PORT) + if env_port and env_port.strip().isdigit(): + return int(env_port.strip()) + + port_file = get_ipc_port_file() + if port_file.exists(): + try: + content = port_file.read_text(encoding="utf-8").strip() + if content.isdigit(): + return int(content) + except Exception as e: + logger.warning(f"Failed to read IPC port from {port_file}: {e}") + + return None + + +def write_ipc_port(port: int) -> Path: + """Save active IPC port to environment and temporary synchronization file.""" + os.environ[ENV_ARTEMIS_IPC_PORT] = str(port) + port_file = get_ipc_port_file() + try: + port_file.parent.mkdir(parents=True, exist_ok=True) + port_file.write_text(str(port), encoding="utf-8") + except Exception as e: + logger.warning(f"Failed to write IPC port file to {port_file}: {e}") + return port_file + + +def clear_ipc_port() -> None: + """Remove IPC port synchronization state file.""" + port_file = get_ipc_port_file() + if port_file.exists(): + try: + port_file.unlink(missing_ok=True) + except Exception as e: + logger.warning(f"Failed to remove IPC port file {port_file}: {e}") + + +def read_ls_address() -> str | None: + """Read the active Language Server address from environment or state file.""" + env_addr = os.getenv(ENV_ANTIGRAVITY_LS_ADDRESS) + if env_addr and env_addr.strip(): + return env_addr.strip() + + addr_file = get_ls_address_file() + if addr_file.exists(): + try: + content = addr_file.read_text(encoding="utf-8").strip() + if content: + return content + except Exception as e: + logger.warning(f"Failed to read LS address from {addr_file}: {e}") + + return None + + +def write_ls_address(address: str) -> Path: + """Save Language Server address to environment and temporary synchronization file.""" + clean_addr = address.strip() + os.environ[ENV_ANTIGRAVITY_LS_ADDRESS] = clean_addr + addr_file = get_ls_address_file() + try: + addr_file.parent.mkdir(parents=True, exist_ok=True) + addr_file.write_text(clean_addr, encoding="utf-8") + except Exception as e: + logger.warning(f"Failed to write LS address file to {addr_file}: {e}") + return addr_file + + +def clear_ls_address() -> None: + """Remove Language Server address synchronization state file.""" + addr_file = get_ls_address_file() + if addr_file.exists(): + try: + addr_file.unlink(missing_ok=True) + except Exception as e: + logger.warning(f"Failed to remove LS address file {addr_file}: {e}") + + +def cleanup_temp_dir(subfolder: str | None = None, max_age_seconds: float | None = None) -> int: + """Clean up expired temporary runtime files in the central temp directory. + + Args: + subfolder: Optional subfolder within the temp directory. + max_age_seconds: Maximum age in seconds before a file is deleted. If None, all files are purged. + + Returns: + Number of files successfully deleted. + """ + temp_dir = get_temp_dir(subfolder) + if not temp_dir.exists(): + return 0 + + now = time.time() + deleted_count = 0 + + for file_path in temp_dir.glob("*"): + if file_path.is_file(): + try: + if max_age_seconds is None or (now - file_path.stat().st_mtime) > max_age_seconds: + file_path.unlink(missing_ok=True) + deleted_count += 1 + except Exception as e: + logger.warning(f"Could not delete temporary file {file_path}: {e}") + + return deleted_count diff --git a/artemis/config/settings.py b/artemis/config/settings.py new file mode 100644 index 0000000..1999909 --- /dev/null +++ b/artemis/config/settings.py @@ -0,0 +1,199 @@ +# 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. + +"""Application settings and environment variable configuration management.""" + +import os +from pathlib import Path + +from dotenv import load_dotenv +from pydantic import Field, SecretStr, model_validator +from pydantic_settings import BaseSettings + +from artemis.config.constants import ( + DEFAULT_ADB_HOST, + DEFAULT_ADB_PORT, + DEFAULT_EXPLORER_VERSION, + DEFAULT_MODEL, + DEFAULT_PROFILE, + ENV_DATA_ENGINE_DB_PATH, + ENV_GCP_API_KEY, + ENV_GEMINI_API_KEY, + ENV_GOOGLE_API_KEY, + ENV_OPEN_ROUTER_API_KEY, + ENV_OPENAI_API_KEY, + ENV_XAI_API_KEY, + ExplorerVersion, +) +from artemis.config.paths import ( + GLOBAL_APP_DIR, + get_app_dir, + get_data_engine_db_path, + get_default_traces_path, + get_temp_dir, +) +from artemis.utils.logger import get_logger + +# Load environment configuration from app dir if present, else current working dir +_global_env = GLOBAL_APP_DIR / ".env" +if _global_env.exists(): + load_dotenv(dotenv_path=_global_env, verbose=True) +else: + load_dotenv(verbose=True) + +logger = get_logger(__name__) + + +class Settings(BaseSettings): + """Centralized ARTEMIS runtime settings loaded from environment and .env files.""" + + # LLM Provider Authentication + OPENAI_API_KEY: SecretStr | None = None + GOOGLE_API_KEY: SecretStr | None = None + GEMINI_API_KEY: SecretStr | None = None + GCP_API_KEY: SecretStr | None = None + ANTHROPIC_API_KEY: SecretStr | None = None + XAI_API_KEY: SecretStr | None = None + OPEN_ROUTER_API_KEY: SecretStr | None = None + + # Google Cloud Vision OCR Authentication + OCR_API_KEY: SecretStr | None = None + VISION_API_KEY: SecretStr | None = None + API_KEY: SecretStr | None = None + + # Custom Provider Endpoints + OPENAI_BASE_URL: str | None = None + + # Android ADB Connectivity + ADB_HOST: str | None = Field(default=DEFAULT_ADB_HOST) + ADB_PORT: int | None = Field(default=DEFAULT_ADB_PORT) + ADB_DEVICE_SERIAL: str | None = None + + # Execution Defaults + PROJECT_NAME: str | None = None + ARTEMIS_DEFAULT_PROFILE: str = Field(default=DEFAULT_PROFILE) + ARTEMIS_DEFAULT_MODEL: str = Field(default=DEFAULT_MODEL) + + # Explorer Tool Settings + EXPLORER_VERSION: ExplorerVersion = Field( + default=DEFAULT_EXPLORER_VERSION, + description="Active version mode for the UI Explorer tool ('flash', 'pro', or 'ultra')", + ) + EXPLORER_CACHING: bool = Field( + default=True, + description="Enable context caching when running multi-turn pro/ultra Explorer", + ) + + # Paths & Storage + TRACES_PATH: Path = Field(default_factory=get_default_traces_path) + DATA_ENGINE_DB_PATH: Path = Field(default_factory=get_data_engine_db_path) + TEMP_PATH: Path = Field(default_factory=get_temp_dir) + + model_config = {"env_file": ".env", "extra": "ignore"} + + @model_validator(mode="after") + def fallback_api_keys(self) -> "Settings": + """Normalize Google / Gemini / GCP and OCR API keys for seamless interoperability.""" + if not self.GOOGLE_API_KEY: + if self.GEMINI_API_KEY: + self.GOOGLE_API_KEY = self.GEMINI_API_KEY + elif self.GCP_API_KEY: + self.GOOGLE_API_KEY = self.GCP_API_KEY + + # Fallback for OCR & Vision API keys + if not self.OCR_API_KEY: + if self.VISION_API_KEY: + self.OCR_API_KEY = self.VISION_API_KEY + elif self.API_KEY: + self.OCR_API_KEY = self.API_KEY + if not self.API_KEY and self.OCR_API_KEY: + self.API_KEY = self.OCR_API_KEY + return self + + def get_api_key(self, provider: str) -> SecretStr | None: + """Get API key for a specified provider. + + Args: + provider: Provider name (e.g., 'google', 'gemini', 'openai', 'ocr', 'vision'). + + Returns: + SecretStr containing the API key or None if not configured. + """ + provider_lower = provider.lower() + if provider_lower in ("google", "gemini", "vertexai"): + return self.GOOGLE_API_KEY or self.GEMINI_API_KEY or self.GCP_API_KEY + elif provider_lower in ("ocr", "vision", "google_vision"): + return self.OCR_API_KEY or self.VISION_API_KEY or self.API_KEY or self.GOOGLE_API_KEY + elif provider_lower == "openai": + return self.OPENAI_API_KEY + elif provider_lower == "anthropic": + return self.ANTHROPIC_API_KEY + elif provider_lower == "openrouter": + return self.OPEN_ROUTER_API_KEY + elif provider_lower == "xai": + return self.XAI_API_KEY + return None + + def set_api_key(self, provider: str, key: str, persist_to_env: bool = False) -> None: + """Dynamically set an API key at runtime, optionally persisting to .env in the app dir. + + Args: + provider: Target provider name. + key: Secret API key string. + persist_to_env: Whether to save the key to the app directory's .env file. + """ + secret = SecretStr(key) + provider_lower = provider.lower() + env_key_name = None + + if provider_lower in ("google", "gemini", "vertexai"): + self.GOOGLE_API_KEY = secret + self.GEMINI_API_KEY = secret + self.GCP_API_KEY = secret + env_key_name = ENV_GOOGLE_API_KEY + os.environ[ENV_GOOGLE_API_KEY] = key + os.environ[ENV_GEMINI_API_KEY] = key + os.environ[ENV_GCP_API_KEY] = key + elif provider_lower == "openai": + self.OPENAI_API_KEY = secret + env_key_name = ENV_OPENAI_API_KEY + os.environ[ENV_OPENAI_API_KEY] = key + elif provider_lower == "openrouter": + self.OPEN_ROUTER_API_KEY = secret + env_key_name = ENV_OPEN_ROUTER_API_KEY + os.environ[ENV_OPEN_ROUTER_API_KEY] = key + elif provider_lower == "xai": + self.XAI_API_KEY = secret + env_key_name = ENV_XAI_API_KEY + os.environ[ENV_XAI_API_KEY] = key + + if persist_to_env and env_key_name: + try: + env_file = get_app_dir() / ".env" + lines = [] + if env_file.exists(): + lines = env_file.read_text(encoding="utf-8").splitlines() + lines = [line for line in lines if not line.startswith(f"{env_key_name}=")] + lines.append(f"{env_key_name}={key}") + env_file.write_text("\n".join(lines) + "\n", encoding="utf-8") + except Exception as e: + logger.warning(f"Could not persist {env_key_name} to .env: {e}") + + +# Singleton instance +settings = Settings() + +# Synchronize DATA_ENGINE_DB_PATH in environment for external sub-processes/tools +if settings.DATA_ENGINE_DB_PATH: + os.environ[ENV_DATA_ENGINE_DB_PATH] = str(settings.DATA_ENGINE_DB_PATH) diff --git a/artemis/constants.py b/artemis/constants.py new file mode 100644 index 0000000..e27b614 --- /dev/null +++ b/artemis/constants.py @@ -0,0 +1,50 @@ +# 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. + +RECURSION_LIMIT = 400 +MAX_MESSAGES_IN_HISTORY = 25 +VALIDATOR_MESSAGES_KEY = "validator_messages" + +# Optimization Constants +CHECKER_MAX_ITERATIONS = 20 +OPERATOR_MAX_CHAT_ROUNDS = 4 +VALIDATOR_POLL_TIMEOUT = 2.0 +VALIDATOR_POLL_INTERVAL = 0.2 +VALIDATOR_OCR_INTERVAL = 1.0 +VALIDATOR_UI_HIERARCHY_TIMEOUT = 1.0 + +from google.genai import types as _genai_types + +SAFETY_SETTINGS_BLOCK_NONE = [ + _genai_types.SafetySetting( + category=_genai_types.HarmCategory.HARM_CATEGORY_HATE_SPEECH, + threshold=_genai_types.HarmBlockThreshold.BLOCK_NONE, + ), + _genai_types.SafetySetting( + category=_genai_types.HarmCategory.HARM_CATEGORY_HARASSMENT, + threshold=_genai_types.HarmBlockThreshold.BLOCK_NONE, + ), + _genai_types.SafetySetting( + category=_genai_types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT, + threshold=_genai_types.HarmBlockThreshold.BLOCK_NONE, + ), + _genai_types.SafetySetting( + category=_genai_types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, + threshold=_genai_types.HarmBlockThreshold.BLOCK_NONE, + ), + _genai_types.SafetySetting( + category=_genai_types.HarmCategory.HARM_CATEGORY_CIVIC_INTEGRITY, + threshold=_genai_types.HarmBlockThreshold.BLOCK_NONE, + ), +] diff --git a/artemis/context.py b/artemis/context.py new file mode 100644 index 0000000..b5ec979 --- /dev/null +++ b/artemis/context.py @@ -0,0 +1,219 @@ +# 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. + +"""Context variables for global state management. + +Uses ContextVar to avoid prop drilling and maintain clean function signatures. +""" + +from __future__ import annotations + +import asyncio + +try: + from enum import StrEnum +except ImportError: + from enum import Enum + + class StrEnum(str, Enum): + pass + + +from pathlib import Path +from typing import Any, Literal, TYPE_CHECKING + +if TYPE_CHECKING: + from artemis.data_engine.engine import DataEngine + +try: + from adbutils import AdbClient +except ImportError: + AdbClient = Any + +from pydantic import BaseModel, ConfigDict, Field, PrivateAttr +from artemis.utils.video import detect_video_tools_enabled + + +from artemis.config import ExplorerConfig, LLMConfig, OutputterConfig + + +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +class AppLaunchResult(BaseModel): + """Result of initial app launch attempt.""" + + model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True) + + locked_app_package: str + locked_app_initial_launch_success: bool | None + locked_app_initial_launch_error: str | None + + +class DevicePlatform(StrEnum): + """Mobile device platform enumeration.""" + + ANDROID = "android" + + +class DeviceContext(BaseModel): + model_config = ConfigDict( + arbitrary_types_allowed=True, + extra="allow", + ) + + host_platform: Literal["WINDOWS", "LINUX"] + mobile_platform: DevicePlatform + device_id: str + + device_width: int + device_height: int + + def to_str(self): + return ( + f"Host platform: {self.host_platform}\n" + f"Mobile platform: {self.mobile_platform.value}\n" + f"Device ID: {self.device_id}\n" + f"Device width: {self.device_width}\n" + f"Device height: {self.device_height}\n" + ) + + +class ExecutionSetup(BaseModel): + """Execution setup for a task.""" + + model_config = ConfigDict( + arbitrary_types_allowed=True, + extra="allow", + ) + + traces_path: Path | None = None + trace_name: str | None = None + enable_remote_tracing: bool = False + app_lock_status: AppLaunchResult | None = None + video_recording_tools_enabled: bool = Field(default_factory=detect_video_tools_enabled) + disable_checker: bool = True + checker_max_iterations: int = 20 + checker_max_chat_rounds: int = 4 + disable_planner_validation: bool = True + planner_validation_threshold: float = 0.85 + enable_committee: bool = False + committee_debate_rounds: int = 2 + disable_outputter: bool = False + outputter: OutputterConfig = Field(default_factory=OutputterConfig) + explorer: ExplorerConfig = Field(default_factory=ExplorerConfig) + explorer_versions: dict[str, str] = Field(default_factory=dict) + + @property + def explorer_version(self) -> str: + return self.explorer.default_version + + @property + def explorer_flash_mode(self) -> str: + return self.explorer.flash_mode + + @property + def explorer_pro_mode(self) -> str: + return self.explorer.pro_mode + + @property + def explorer_caching(self) -> bool: + return self.explorer.caching + + def get_locked_app_package(self) -> str | None: + """Get the locked app package name if app locking is enabled. + + Returns: + The locked app package name, or None if app locking is not enabled. + """ + if self.app_lock_status: + return self.app_lock_status.locked_app_package + return None + + +class ArtemisContext(BaseModel): + model_config = ConfigDict( + arbitrary_types_allowed=True, + extra="allow", + ) + + device: DeviceContext + llm_config: LLMConfig | None = None + model_router: Any | None = None + agent_config: Any = None + adb_client: Any | None = None + ui_adb_client: Any | None = None + + execution_setup: ExecutionSetup | None = None + + data_engine: DataEngine | None = None + checker_task: asyncio.Task[Any] | None = None + planner_task: asyncio.Task[Any] | None = None + background_tasks: list[asyncio.Task[Any]] = Field(default_factory=list) + package_cache: dict[str, str | None] = Field(default_factory=dict) + background_jobs: dict[str, dict[str, Any]] = Field(default_factory=dict) + + task_plan_snapshot: Path | None = None + task_plan_content_before: str | None = None + + _genai_client: Any | None = PrivateAttr(default=None) + mcp_client_ctx: Any | None = None + mcp_session: Any | None = None + + async def __aenter__(self) -> ArtemisContext: + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + if self.background_tasks: + logger.info(f"Waiting for {len(self.background_tasks)} background tasks to complete...") + results = await asyncio.gather(*self.background_tasks, return_exceptions=True) + for task, result in zip(self.background_tasks, results): + if isinstance(result, Exception): + logger.error( + f"Background task {task.get_name()} failed: {result}", + exc_info=result, + ) + self.background_tasks.clear() + + if self.mcp_client_ctx: + try: + await self.mcp_client_ctx.__aexit__(exc_type, exc_val, exc_tb) + except Exception: + pass + finally: + self.mcp_client_ctx = None + self.mcp_session = None + + if self.data_engine: + try: + await self.data_engine.shutdown() + except Exception: + pass + + def get_adb_client(self) -> Any: + if self.adb_client is None: + raise ValueError("No ADB client in context.") + return self.adb_client + + def get_ui_adb_client(self) -> Any: + if self.ui_adb_client is None: + raise ValueError("No UIAutomator client in context.") + return self.ui_adb_client + + +from artemis.data_engine.engine import DataEngine + +ArtemisContext.model_rebuild() diff --git a/artemis/controllers/__init__.py b/artemis/controllers/__init__.py new file mode 100644 index 0000000..0d3e144 --- /dev/null +++ b/artemis/controllers/__init__.py @@ -0,0 +1,41 @@ +# 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. + +"""Backward compatibility layer for legacy artemis.controllers module. +All new development should use artemis.drivers instead. +""" + +from artemis.drivers.base import BaseDeviceDriver, KeyCode, ScreenData, SwipeDirection +from artemis.drivers.factory import create_driver, get_driver +from artemis.drivers.types import ( + Bounds, + CoordinatesSelectorRequest, + PercentagesSelectorRequest, + SwipeRequest, + TapOutput, +) + +__all__ = [ + "BaseDeviceDriver", + "KeyCode", + "ScreenData", + "SwipeDirection", + "Bounds", + "CoordinatesSelectorRequest", + "PercentagesSelectorRequest", + "SwipeRequest", + "TapOutput", + "create_driver", + "get_driver", +] diff --git a/artemis/controllers/android_controller.py b/artemis/controllers/android_controller.py new file mode 100644 index 0000000..53b0395 --- /dev/null +++ b/artemis/controllers/android_controller.py @@ -0,0 +1,912 @@ +# 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 asyncio +import base64 +from io import BytesIO +from pathlib import Path +import re +import tempfile +import time +from typing import Any +from uuid import uuid4 + +from adbutils import AdbClient, AdbDevice +from PIL import Image + +from artemis.clients.ui_automator_client import ( + UIAutomatorClient, + _parse_hierarchy_xml_to_elements, +) +from artemis.config import get_temp_dir +from artemis.controllers.device_controller import ( + MobileDeviceController, + ScreenDataResponse, +) +from artemis.controllers.types import ( + Bounds, + CoordinatesSelectorRequest, + TapOutput, +) +from artemis.data_engine.models import VideoRecordingRecord +from artemis.utils.logger import get_logger +from artemis.utils.video import ( + DEFAULT_MAX_DURATION_SECONDS, + RecordingSession, + VideoRecordingResult, + cleanup_video_segments, + concatenate_videos, + get_active_session, + has_active_session, + remove_active_session, + set_active_session, + trim_video, +) + +logger = get_logger(__name__) + + +class AndroidDeviceController(MobileDeviceController): + def __init__( + self, + device_id: str, + adb_client: AdbClient, + ui_adb_client: UIAutomatorClient, + device_width: int, + device_height: int, + data_engine_start_time: float | None = None, + data_engine: Any | None = None, + ): + self.device_id = device_id + self.adb_client = adb_client + self.ui_adb_client = ui_adb_client + self.device_width = device_width + self.device_height = device_height + self.data_engine_start_time = data_engine_start_time + self.data_engine = data_engine + self._device: AdbDevice | None = None + self._segment_cache: dict[tuple[float, float | None], VideoRecordingResult] = {} + + @property + def device(self) -> AdbDevice: + if self._device is None: + self._device = self.adb_client.device(serial=self.device_id) + return self._device + + async def tap( + self, + coords: CoordinatesSelectorRequest, + long_press: bool = False, + long_press_duration: int = 1000, + times: int = 1, + delay_ms: int = 100, + ) -> TapOutput: + try: + if long_press: + cmd = ( + f"input swipe {coords.x} {coords.y} {coords.x} {coords.y} {long_press_duration}" + ) + else: + if times <= 1: + cmd = f"input tap {coords.x} {coords.y}" + else: + # Chain multiple taps in a single shell command to eliminate shell spawning overhead + # between rapid consecutive clicks (e.g. 7 taps to enter developer mode) + taps = [f"input tap {coords.x} {coords.y}"] * times + cmd = f" && sleep {delay_ms / 1000.0:.3f} && ".join(taps) + + logger.info(f"Executing ADB shell: {cmd}") + self.device.shell(cmd) + return TapOutput(error=None) + except Exception as e: + return TapOutput(error=f"ADB tap failed: {str(e)}") + + async def swipe( + self, + start: CoordinatesSelectorRequest, + end: CoordinatesSelectorRequest, + duration: int = 400, + ) -> str | None: + try: + cmd = f"input touchscreen swipe {start.x} {start.y} {end.x} {end.y} {duration}" + logger.info(f"Executing ADB command: {cmd}") + self.device.shell(cmd) + return None + except Exception as e: + return f"ADB swipe failed: {str(e)}" + + async def get_screen_data(self) -> ScreenDataResponse: + """Get screen data using the UIAutomator2 client""" + try: + logger.info("Using UIAutomator2 for screen data retrieval") + ui_data = self.ui_adb_client.get_screen_data() + return ScreenDataResponse( + base64=ui_data.base64, + elements=ui_data.elements, + width=ui_data.width, + height=ui_data.height, + platform="android", + ) + except Exception as e: + logger.error(f"Failed to get screen data: {e}") + raise + + async def screenshot(self) -> str: + try: + b64 = self.ui_adb_client.get_screenshot_base64() + if b64 is None: + raise RuntimeError("Failed to capture screenshot") + return b64 + except Exception as e: + logger.error(f"Failed to take screenshot: {e}") + raise + + async def input_text(self, text: str) -> bool: + try: + # Fast path for ASCII text: Use ADB directly. + # This avoids toggling FastInputIME, which hides the keyboard and can cause focus loss. + is_ascii = all(ord(c) < 128 for c in text) + if is_ascii: + return self._input_text_adb_fallback(text) + + self.ui_adb_client.send_text(text) + return True + except Exception as e: + logger.warning(f"UIAutomator2 send_text failed: {e}, falling back to ADB shell") + return self._input_text_adb_fallback(text) + + def _input_text_adb_fallback(self, text: str) -> bool: + """Fallback method using ADB shell input text command.""" + try: + lines = text.split("\n") + for line_idx, line in enumerate(lines): + if line_idx > 0: + self.device.shell("input keyevent 66") + if not line: + continue + parts = line.split("%s") + for i, part in enumerate(parts): + to_write = "" + for char in part: + if char == " ": + to_write += "%s" + elif char in [ + "&", + "<", + ">", + "|", + ";", + "(", + ")", + "$", + "`", + "\\", + '"', + "'", + ]: + to_write += f"\\{char}" + else: + to_write += char + + if to_write: + self.device.shell(f"input text '{to_write}'") + + if i < len(parts) - 1: + self.device.shell("input keyevent 62") + + return True + except Exception as e: + logger.error(f"Failed to input text via ADB fallback: {e}") + return False + + async def launch_app(self, package_or_bundle_id: str) -> bool: + try: + # Clear foreground obstructions before launching + logger.info( + "Clearing foreground obstructions (status bar, system dialogs," + " keyboard) before launching app..." + ) + try: + # 1. Collapse status bar / notification shade + self.device.shell("cmd statusbar collapse") + except Exception as e: + logger.warning(f"Failed to collapse status bar: {e}") + + try: + # 2. Close system dialogs + self.device.shell("am broadcast -a android.intent.action.CLOSE_SYSTEM_DIALOGS") + except Exception as e: + logger.warning(f"Failed to close system dialogs: {e}") + + # 3. Dismiss keyboard if visible + if self.is_keyboard_visible(): + await self.dismiss_keyboard() + + self.device.app_start(package_or_bundle_id) + return True + except Exception as e: + logger.error(f"Failed to launch app {package_or_bundle_id}: {e}") + return False + + async def terminate_app(self, package_or_bundle_id: str | None) -> bool: + try: + if package_or_bundle_id is None: + current_app = self._get_current_foreground_package() + if current_app: + logger.info(f"Stopping currently running app: {current_app}") + self.device.app_stop(current_app) + else: + logger.warning("No foreground app detected") + return False + else: + self.device.app_stop(package_or_bundle_id) + return True + except Exception as e: + logger.error(f"Failed to terminate app {package_or_bundle_id}: {e}") + return False + + async def open_url(self, url: str) -> bool: + try: + cmd = f"am start -a android.intent.action.VIEW -d {url}" + logger.info(f"Executing ADB shell: {cmd}") + self.device.shell(cmd) + return True + except Exception as e: + logger.error(f"Failed to open URL {url}: {e}") + return False + + def is_keyboard_visible(self) -> bool: + try: + res = self.device.shell("dumpsys input_method") + if isinstance(res, bytes): + res = res.decode("utf-8") + + # Common patterns for open keyboard in dumpsys + patterns = [ + "mInputShown=true", + "mInputViewShowing=true", + "mInputViewShowing=VISIBLE", + ] + for p in patterns: + if p in res: + return True + + # Alternate check using dumpsys window + res_win = self.device.shell( + "dumpsys window | grep -E 'mInputMethodWindow|mShowingInputMethod'" + ) + if isinstance(res_win, bytes): + res_win = res_win.decode("utf-8") + if "mInputMethodWindow" in res_win and "visible=true" in res_win.lower(): + return True + + return False + except Exception as e: + logger.error(f"Failed to check keyboard visibility: {e}") + return False + + async def dismiss_keyboard(self) -> bool: + try: + if not self.is_keyboard_visible(): + return True + + logger.info("Keyboard is visible, attempting to dismiss via Keyevent 111 (Escape)") + self.device.shell("input keyevent 111") + await asyncio.sleep(0.3) + + if not self.is_keyboard_visible(): + logger.info("Keyboard successfully dismissed via Keyevent 111") + return True + + logger.info("Keyboard still visible, trying Keyevent 4 (BACK) as fallback to dismiss") + self.device.shell("input keyevent 4") + await asyncio.sleep(0.3) + + if not self.is_keyboard_visible(): + logger.info("Keyboard successfully dismissed via Keyevent 4") + return True + + logger.warning("Failed to dismiss keyboard after Escape and BACK") + return False + except Exception as e: + logger.error(f"Error in dismiss_keyboard: {e}") + return False + + async def press_back(self) -> bool: + try: + if self.is_keyboard_visible(): + logger.info( + "Intercepted press_back: Keyboard is visible. Attempting to" + " dismiss keyboard first." + ) + return await self.dismiss_keyboard() + + logger.info("Executing normal press_back (keyboard not visible)") + self.device.shell("input keyevent 4") + return True + except Exception as e: + logger.error(f"Failed to press back: {e}") + return False + + async def press_home(self) -> bool: + try: + self.device.shell("input keyevent 3") + return True + except Exception as e: + logger.error(f"Failed to press home: {e}") + return False + + async def press_enter(self) -> bool: + try: + self.device.shell("input keyevent 66") + return True + except Exception as e: + logger.error(f"Failed to press enter: {e}") + return False + + async def press_key(self, keycode: str) -> bool: + try: + # Intercept BACK keycodes to handle keyboard dismissal first + clean_keycode = str(keycode).strip().upper() + if clean_keycode in ("4", "KEYCODE_BACK", "BACK"): + if self.is_keyboard_visible(): + logger.info( + f"Intercepted press_key({keycode}): Keyboard is" + " visible. Dismissing keyboard first." + ) + return await self.dismiss_keyboard() + + self.device.shell(f"input keyevent {keycode}") + return True + except Exception as e: + logger.error(f"Failed to press key {keycode}: {e}") + return False + + async def get_ui_hierarchy(self) -> list[dict]: + + max_attempts = 4 + retry_delay = 0.5 + + for attempt in range(1, max_attempts + 1): + try: + hierarchy_xml = self.ui_adb_client.get_hierarchy() + elements = _parse_hierarchy_xml_to_elements(hierarchy_xml) + + # Check if hierarchy contains actual UI elements besides the root node + if len(elements) > 1: + if attempt > 1: + logger.info( + f"Successfully retrieved non-empty UI hierarchy on attempt {attempt}" + ) + return elements + + logger.warning( + "Retrieved empty or root-only UI hierarchy (length" + f" {len(elements)}) on attempt {attempt}/{max_attempts}." + f" Retrying in {retry_delay}s..." + ) + except Exception as e: + logger.warning( + f"Failed to get UI hierarchy on attempt {attempt}/{max_attempts}: {e}" + ) + if attempt == max_attempts: + logger.error(f"Failed to get UI hierarchy after {max_attempts} attempts: {e}") + return [] + + if attempt < max_attempts: + await asyncio.sleep(retry_delay) + + return [] + + def find_element( + self, + ui_hierarchy: list[dict], + resource_id: str | None = None, + text: str | None = None, + index: int = 0, + ) -> tuple[dict | None, Bounds | None, str | None]: + if not resource_id and not text: + return None, None, "No resource_id or text provided" + + matches = [] + for element in ui_hierarchy: + if resource_id and element.get("resource-id") == resource_id: + matches.append(element) + elif text and (element.get("text") == text or element.get("accessibilityText") == text): + matches.append(element) + + if not matches: + criteria = f"resource_id='{resource_id}'" if resource_id else f"text='{text}'" + return None, None, f"No element found with {criteria}" + + if index >= len(matches): + criteria = f"resource_id='{resource_id}'" if resource_id else f"text='{text}'" + return ( + None, + None, + (f"Index {index} out of range for {criteria} (found {len(matches)} matches)"), + ) + + element = matches[index] + bounds = self._extract_bounds(element) + + return element, bounds, None + + def _get_current_foreground_package(self) -> str | None: + try: + app_info = self.device.current_app() + if app_info and app_info.package: + return app_info.package + return None + except Exception as e: + logger.error( + f"Failed to get current foreground package using adbutils: {e}." + " Falling back to dumpsys parsing." + ) + try: + result = self.device.shell("dumpsys window | grep mCurrentFocus") + + # Convert to string if bytes + if isinstance(result, bytes): + result_str = result.decode("utf-8") + elif isinstance(result, str): + result_str = result + else: + return None + + if result_str and "=" in result_str: + parts = result_str.split("/") + if len(parts) > 0: + package = parts[0].split()[-1] + return package if package else None + except Exception as fallback_err: + logger.error(f"Fallback dumpsys parsing also failed: {fallback_err}") + return None + + def _extract_bounds(self, element: dict) -> Bounds | None: + bounds_str = element.get("bounds") + if not bounds_str or not isinstance(bounds_str, str): + return None + + try: + match = re.match(r"\[(\d+),(\d+)\]\[(\d+),(\d+)\]", bounds_str) + if match: + return Bounds( + x1=int(match.group(1)), + y1=int(match.group(2)), + x2=int(match.group(3)), + y2=int(match.group(4)), + ) + except (ValueError, IndexError): + return None + + return None + + async def erase_text(self, nb_chars: int | None = None) -> bool: + # 1. If clearing the whole field, try fast UIAutomator2 clear first + if nb_chars is None: + logger.info("Attempting fast UIAutomator2 clear_text...") + if self.ui_adb_client.clear_text(): + return True + logger.warning("UIAutomator2 clear_text failed, falling back to optimized ADB shell") + + # 2. Optimized ADB fallback: combine multiple KEYCODE_DEL keyevents into a single command + try: + chars_to_delete = nb_chars if nb_chars is not None else 50 + if chars_to_delete > 0: + # KEYCODE_DEL is 67. Build arguments list "67 67 67 ..." + keycode_args = " ".join(["67"] * chars_to_delete) + cmd = f"input keyevent {keycode_args}" + logger.info(f"Executing combined ADB shell delete: {cmd}") + self.device.shell(cmd) + return True + except Exception as e: + logger.error(f"Failed to erase text via ADB fallback: {e}") + return False + + async def cleanup(self) -> None: + pass + + def get_compressed_b64_screenshot(self, image_base64: str, quality: int = 50) -> str: + if image_base64.startswith("data:image"): + image_base64 = image_base64.split(",")[1] + + image_data = base64.b64decode(image_base64) + image = Image.open(BytesIO(image_data)) + + compressed_io = BytesIO() + image.save(compressed_io, format="JPEG", quality=quality, optimize=True) + + compressed_base64 = base64.b64encode(compressed_io.getvalue()).decode("utf-8") + return compressed_base64 + + async def start_video_recording( + self, + max_duration_seconds: int = DEFAULT_MAX_DURATION_SECONDS, + output_dir: Path | None = None, + ) -> VideoRecordingResult: + """Start screen recording on Android device using scrcpy.""" + self._segment_cache.clear() + if has_active_session(self.device_id): + return VideoRecordingResult( + success=False, + message=(f"Recording already in progress for device {self.device_id}"), + ) + + try: + # Create a temp file for the MKV recording if output_dir not provided + if not output_dir: + output_dir = Path( + tempfile.mkdtemp(prefix="scrcpy_", dir=get_temp_dir("recordings")) + ) + else: + output_dir.mkdir(parents=True, exist_ok=True) + local_video_path = output_dir / "recording.mkv" + + video_id = uuid4() + session = RecordingSession( + video_id=video_id, + device_id=self.device_id, + start_time=time.time(), + data_engine_start_time=self.data_engine_start_time, + local_video_path=local_video_path, + ) + + # Persist to local database if Data Engine is active + if self.data_engine and self.data_engine.storage: + try: + record = VideoRecordingRecord( + video_id=video_id, + session_id=self.data_engine.current_session_id, + device_id=self.device_id, + start_time=session.start_time, + local_video_path=str(local_video_path), + ) + self.data_engine.storage.create_video_recording(record) + except Exception as db_err: + logger.error(f"Failed to persist video recording start to DB: {db_err}") + + # Start scrcpy in background + cmd = [ + "scrcpy", + "--serial", + self.device_id, + "--no-window", + "--record", + str(local_video_path), + "--record-format", + "mkv", + "--video-bit-rate", + "2M", + ] + + process = await asyncio.create_subprocess_exec( + *cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + + session.process = process + set_active_session(self.device_id, session) + + logger.info( + f"Started scrcpy recording on {self.device_id}, saving to {local_video_path}" + ) + + # Wait a bit and check if it's still running. + await asyncio.sleep(1.0) + if process.returncode is not None: + stderr = await process.stderr.read() + logger.error(f"scrcpy failed to start: {stderr.decode()}") + remove_active_session(self.device_id) + return VideoRecordingResult( + success=False, + message=f"scrcpy failed to start: {stderr.decode()}", + ) + + return VideoRecordingResult( + success=True, + message=f"Recording started on {self.device_id}", + ) + + except Exception as e: + logger.error(f"Failed to start scrcpy recording: {e}") + remove_active_session(self.device_id) + return VideoRecordingResult( + success=False, + message=f"Failed to start recording: {e}", + ) + + async def stop_video_recording(self) -> VideoRecordingResult: + """Stop scrcpy recording and return the video file.""" + self._segment_cache.clear() + session = get_active_session(self.device_id) + if not session: + return VideoRecordingResult( + success=False, + message=f"No active recording for device {self.device_id}", + ) + + try: + process = session.process + if process is not None: + process.terminate() + await asyncio.wait_for(process.wait(), timeout=5.0) + + output_path = session.local_video_path + + # Fix: Concatenate previous segments if scrcpy crashed and restarted + if session.android_video_segments: + all_segments = session.android_video_segments + [output_path] + concatenated_mkv = output_path.parent / "full_recording.mkv" + + if await concatenate_videos(all_segments, concatenated_mkv): + cleanup_video_segments(all_segments) + output_path = concatenated_mkv + + # Convert MKV to MP4 + mp4_path = ( + output_path.parent / "recording.mp4" + if output_path.name == "full_recording.mkv" + else output_path.with_suffix(".mp4") + ) + success = await self._convert_mkv_to_mp4(output_path, mp4_path) + + remove_active_session(self.device_id) + + final_video_path = mp4_path if (success and mp4_path.exists()) else output_path + + # Persist update to local database if Data Engine is active + if self.data_engine and self.data_engine.storage: + try: + record = VideoRecordingRecord( + video_id=session.video_id, + session_id=self.data_engine.current_session_id, + device_id=self.device_id, + start_time=session.start_time, + end_time=time.time(), + local_video_path=str(final_video_path), + ) + self.data_engine.storage.update_video_recording(record) + except Exception as db_err: + logger.error(f"Failed to persist video recording stop to DB: {db_err}") + + if success and mp4_path.exists(): + try: + output_path.unlink() + except Exception: + pass + return VideoRecordingResult( + success=True, + message=f"Recording stopped, saved to {mp4_path}", + video_path=mp4_path, + ) + else: + return VideoRecordingResult( + success=True, + message=( + f"Recording stopped, saved as MKV to {output_path} (conversion failed)" + ), + video_path=output_path, + ) + + except Exception as e: + logger.error(f"Failed to stop scrcpy recording: {e}") + remove_active_session(self.device_id) + return VideoRecordingResult( + success=False, + message=f"Failed to stop recording: {e}", + ) + + async def extract_segment_metadata( + self, + start_relative_time: float, + end_relative_time: float | None = None, + ) -> VideoRecordingResult: + """Get a video segment for a specific time range (relative to video start).""" + cache_key = ( + round(start_relative_time, 1), + round(end_relative_time, 1) if end_relative_time is not None else None, + ) + if cache_key in self._segment_cache: + cached_res = self._segment_cache[cache_key] + if cached_res.success and cached_res.video_path and cached_res.video_path.exists(): + logger.info( + "Reusing cached trimmed video segment for range" + f" {cache_key[0]}s to {cache_key[1]}s" + ) + return cached_res + + session = get_active_session(self.device_id) + if not session: + return VideoRecordingResult( + success=False, + message=f"No active recording for device {self.device_id}", + ) + + try: + mkv_path = session.local_video_path + if not mkv_path or not mkv_path.exists(): + return VideoRecordingResult( + success=False, + message="Recording file not found", + ) + + # Check if process is still running (best effort for disconnection check) + crashed = False + restart_message = "" + + if session.process and session.process.returncode is not None: + logger.warning(f"scrcpy process for {self.device_id} has terminated unexpectedly") + crashed = True + + # Move current file to segments list (save it before overwriting session state) + session.android_video_segments.append(mkv_path) + + # Try to restart scrcpy + try: + session.android_segment_index += 1 + output_dir = mkv_path.parent + new_video_path = output_dir / f"recording_{session.android_segment_index}.mkv" + + cmd = [ + "scrcpy", + "--serial", + self.device_id, + "--no-window", + "--record", + str(new_video_path), + "--record-format", + "mkv", + "--video-bit-rate", + "2M", + ] + + process = await asyncio.create_subprocess_exec( + *cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + + # Wait a bit and check if it's still running. + await asyncio.sleep(1.0) + if process.returncode is None: + session.process = process + session.local_video_path = new_video_path + restart_message = "Successfully restarted scrcpy." + else: + stderr = await process.stderr.read() + restart_message = f"Failed to restart scrcpy: {stderr.decode()}" + except Exception as e: + restart_message = f"Failed to restart scrcpy: {e}" + + current_time = time.time() + video_duration = current_time - session.start_time + + # Calculate offset + offset = 0.0 + if session.data_engine_start_time is not None: + offset = session.start_time - session.data_engine_start_time + if offset < 0: + offset = 0.0 + + video_start_relative_time = start_relative_time - offset + video_end_relative_time = ( + end_relative_time - offset if end_relative_time is not None else None + ) + + # Safety buffer: avoid reading the very last second being written + safe_duration = max(0.0, video_duration - 1.0) + + truncation_warning = None + if video_end_relative_time is not None and video_end_relative_time > safe_duration: + truncation_warning = ( + "Video segment truncated. Requested end time" + f" {end_relative_time:.1f}s exceeded available safe" + f" duration. Truncated to {safe_duration + offset:.1f}s." + ) + + if video_end_relative_time is None or video_end_relative_time > safe_duration: + video_end_relative_time = safe_duration + + if video_start_relative_time < 0: + video_start_relative_time = 0.0 + + if video_start_relative_time >= video_end_relative_time: + return VideoRecordingResult( + success=False, + message=( + "Invalid time range or requested too close to current" + f" time. Available duration: {safe_duration:.1f}s" + ), + ) + + # Trim directly from MKV to MP4 (using mkv_path which points to the file before we potentially changed it for restart) + + trim_output_dir = tempfile.mkdtemp( + prefix="video_trimmed_", + dir=get_temp_dir("trimmed_videos"), + ) + trim_output_path = Path(trim_output_dir) / "segment.mp4" + + # Retry mechanism for ffmpeg stability + success = False + for attempt in range(3): + success = await trim_video( + mkv_path, + video_start_relative_time, + video_end_relative_time, + trim_output_path, + ) + if success and trim_output_path.exists(): + break + logger.warning(f"ffmpeg trim attempt {attempt + 1} failed, retrying in 0.5s...") + await asyncio.sleep(0.5) + + if not success or not trim_output_path.exists(): + return VideoRecordingResult( + success=False, + message="Failed to trim video segment after retries", + ) + + actual_start = video_start_relative_time + offset + actual_end = video_end_relative_time + offset + + message = f"Video segment retrieved for range {actual_start:.1f}s to {actual_end:.1f}s" + if crashed: + message += ( + f". Warning: scrcpy crashed. Data loss may have occurred. {restart_message}" + ) + + file_size_mb = trim_output_path.stat().st_size / (1024 * 1024) + duration = video_end_relative_time - video_start_relative_time + + res = VideoRecordingResult( + success=True, + message=message, + video_path=trim_output_path, + file_size_mb=round(file_size_mb, 2), + duration_seconds=round(duration, 2), + actual_start_relative_time=actual_start, + warning=truncation_warning, + ) + self._segment_cache[cache_key] = res + return res + + except Exception as e: + logger.error(f"Failed to get video segment: {e}") + return VideoRecordingResult( + success=False, + message=f"Failed to get video segment: {e}", + ) + + async def _convert_mkv_to_mp4(self, mkv_path: Path, mp4_path: Path) -> bool: + """Convert MKV to MP4 using ffmpeg.""" + try: + process = await asyncio.create_subprocess_exec( + "ffmpeg", + "-y", + "-i", + str(mkv_path), + "-c", + "copy", + str(mp4_path), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + await process.wait() + return mp4_path.exists() + except Exception as e: + logger.error(f"Failed to convert MKV to MP4: {e}") + return False diff --git a/artemis/controllers/controller_factory.py b/artemis/controllers/controller_factory.py new file mode 100644 index 0000000..acbb486 --- /dev/null +++ b/artemis/controllers/controller_factory.py @@ -0,0 +1,29 @@ +# 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. + +from artemis.context import ArtemisContext +from artemis.controllers.unified_controller import UnifiedMobileController +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +def create_device_controller(ctx: ArtemisContext) -> UnifiedMobileController: + """Factory creating the unified mobile controller delegating to device drivers.""" + return UnifiedMobileController(ctx) + + +def get_controller(ctx: ArtemisContext) -> UnifiedMobileController: + """Backward-compatible alias for create_device_controller.""" + return create_device_controller(ctx) diff --git a/artemis/controllers/device_controller.py b/artemis/controllers/device_controller.py new file mode 100644 index 0000000..d5c2457 --- /dev/null +++ b/artemis/controllers/device_controller.py @@ -0,0 +1,245 @@ +# 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. + +from abc import ABC, abstractmethod +from typing import Any + +from artemis.controllers.types import Bounds, CoordinatesSelectorRequest, TapOutput +from artemis.utils.video import VideoRecordingResult +from pydantic import BaseModel, Field + + +class _CyFunctionDetectorMeta(type): + def __instancecheck__(self, instance): + name = type(instance).__name__ + return ( + name + in ( + "cyfunction", + "cython_function_or_method", + "builtin_function_or_method", + ) + or "cyfunction" in name.lower() + ) + + +class CyFunctionDetector(metaclass=_CyFunctionDetectorMeta): + pass + + +class ScreenDataResponse(BaseModel): + model_config = {"ignored_types": (CyFunctionDetector,)} + base64: str = Field(default="", description="Base64 encoded screenshot string") + elements: list[dict[str, Any]] = Field( + default_factory=list, description="Parsed UI hierarchy elements" + ) + width: int = Field(default=1080, description="Device screen width in pixels") + height: int = Field(default=2400, description="Device screen height in pixels") + platform: str = Field(default="android", description="Platform identifier") + + +class MobileDeviceController(ABC): + @abstractmethod + async def tap( + self, + coords: CoordinatesSelectorRequest, + long_press: bool = False, + long_press_duration: int = 1000, + times: int = 1, + delay_ms: int = 100, + ) -> TapOutput: + raise NotImplementedError("Subclasses must implement this method") + + @abstractmethod + async def swipe( + self, + start: CoordinatesSelectorRequest, + end: CoordinatesSelectorRequest, + duration: int = 400, + ) -> str | None: + """Swipe from start to end coordinates. + + Returns error message on failure, None on success. + """ + raise NotImplementedError("Subclasses must implement this method") + + @abstractmethod + async def screenshot(self) -> str: + """Take a screenshot and return raw image data.""" + raise NotImplementedError("Subclasses must implement this method") + + @abstractmethod + async def input_text(self, text: str) -> bool: + """Input text at the currently focused field. + + Returns True on success, False on failure. + """ + raise NotImplementedError("Subclasses must implement this method") + + @abstractmethod + async def launch_app(self, package_or_bundle_id: str) -> bool: + """Launch an application by package name (Android). + + Returns True on success, False on failure. + """ + raise NotImplementedError("Subclasses must implement this method") + + @abstractmethod + async def terminate_app(self, package_or_bundle_id: str | None) -> bool: + """Terminate an application. + + Returns True on success, False on failure. + """ + raise NotImplementedError("Subclasses must implement this method") + + @abstractmethod + async def open_url(self, url: str) -> bool: + """Open a URL. + + Returns True on success, False on failure. + """ + raise NotImplementedError("Subclasses must implement this method") + + @abstractmethod + async def press_back(self) -> bool: + """Press the back button (Android). + + Returns True on success, False on failure. + """ + raise NotImplementedError("Subclasses must implement this method") + + @abstractmethod + async def press_home(self) -> bool: + """Press the home button. + + Returns True on success, False on failure. + """ + raise NotImplementedError("Subclasses must implement this method") + + @abstractmethod + async def press_enter(self) -> bool: + """Press the enter/return key. + + Returns True on success, False on failure. + """ + raise NotImplementedError("Subclasses must implement this method") + + @abstractmethod + async def press_key(self, keycode: str) -> bool: + """Press a specific key event. + + Returns True on success, False on failure. + """ + raise NotImplementedError("Subclasses must implement this method") + + @abstractmethod + async def get_ui_hierarchy(self) -> list[dict]: + """Get the UI element hierarchy. + + Returns a list of UI elements with their properties. + """ + raise NotImplementedError("Subclasses must implement this method") + + @abstractmethod + def find_element( + self, + ui_hierarchy: list[dict], + resource_id: str | None = None, + text: str | None = None, + index: int = 0, + ) -> tuple[dict | None, Bounds | None, str | None]: + """Find a UI element in the hierarchy. + + Returns: + Tuple of (element_dict, bounds, error_message) + """ + raise NotImplementedError("Subclasses must implement this method") + + @abstractmethod + async def cleanup(self) -> None: + """Clean up resources (e.g., stop companion processes).""" + raise NotImplementedError("Subclasses must implement this method") + + @abstractmethod + async def erase_text(self, nb_chars: int | None = None) -> bool: + """Erase the last nb_chars characters. + + Returns True on success, False on failure. + """ + raise NotImplementedError("Subclasses must implement this method") + + @abstractmethod + async def get_screen_data(self) -> "ScreenDataResponse": + """Get screen data including screenshot (base64), UI hierarchy elements, + + screen dimensions, and platform. + + Returns: + ScreenDataResponse with base64 screenshot, elements, width, height, + platform + """ + raise NotImplementedError("Subclasses must implement this method") + + @abstractmethod + def get_compressed_b64_screenshot(self, image_base64: str, quality: int = 50) -> str: + """Compress a base64 image. + + Returns the compressed base64 image. + """ + raise NotImplementedError("Subclasses must implement this method") + + @abstractmethod + async def start_video_recording( + self, + max_duration_seconds: int = 900, + ) -> VideoRecordingResult: + """Start screen recording on the device. + + Args: + max_duration_seconds: Maximum recording duration in seconds. + + Returns: + VideoRecordingResult with success status and message. + """ + raise NotImplementedError("Subclasses must implement this method") + + @abstractmethod + async def stop_video_recording(self) -> VideoRecordingResult: + """Stop screen recording and retrieve the video file. + + Returns: + VideoRecordingResult with success status, message, and video_path if + successful. + """ + raise NotImplementedError("Subclasses must implement this method") + + @abstractmethod + async def extract_segment_metadata( + self, + start_relative_time: float, + end_relative_time: float | None = None, + ) -> VideoRecordingResult: + """Get a video segment for a specific time range. + + Args: + start_relative_time: Start time in seconds relative to session + start. + end_relative_time: End time in seconds relative to session start, or + None for latest. + + Returns: + VideoRecordingResult with success status, message, and video_path if + successful. + """ + raise NotImplementedError("Subclasses must implement this method") diff --git a/artemis/controllers/platform_specific_commands_controller.py b/artemis/controllers/platform_specific_commands_controller.py new file mode 100644 index 0000000..fd08af4 --- /dev/null +++ b/artemis/controllers/platform_specific_commands_controller.py @@ -0,0 +1,137 @@ +# 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 os +from shutil import which +import time +from typing import Any + +from artemis.context import ArtemisContext, DevicePlatform +from artemis.utils.logger import ArtemisLogger, get_logger +from artemis.utils.shell_utils import run_shell_command_on_host + +logger = get_logger(__name__) + + +def get_adb_device(ctx: ArtemisContext) -> Any: + """Retrieve the ADB device instance (AdbDevice locally or RemoteAdbDevice in Cloud mode).""" + if ctx.adb_client is None: + if os.environ.get("ARTEMIS_CLOUD_MODE") == "1": + from cloud_service.virtualization import RemoteAdbClient + + ctx.adb_client = RemoteAdbClient() + return ctx.adb_client.device(serial=ctx.device.device_id) + return None + + if ctx.device.mobile_platform != DevicePlatform.ANDROID: + return None + + adb = ctx.get_adb_client() + return adb.device(serial=ctx.device.device_id) + + +def get_first_device( + logger: ArtemisLogger | None = None, +) -> tuple[str | None, DevicePlatform | None, None]: + """Gets the first available device.""" + if which("adb"): + try: + android_output = run_shell_command_on_host("adb devices") + lines = android_output.strip().split("\n") + for line in lines: + if "device" in line and not line.startswith("List of devices"): + return line.split()[0], DevicePlatform.ANDROID, None + except RuntimeError as e: + if logger: + logger.error(f"ADB command failed: {e}") + + return None, None, None + + +def get_device_date(ctx: ArtemisContext) -> str: + device = get_adb_device(ctx) + if not device: + return time.strftime("%Y-%m-%d %H:%M:%S") + try: + return str(device.shell("date")).strip() + except Exception as e: + logger.debug(f"Failed to get device date: {e}") + return time.strftime("%Y-%m-%d %H:%M:%S") + + +def list_packages(ctx: ArtemisContext) -> str: + """List installed packages dynamically from the physical/remote device.""" + device = get_adb_device(ctx) + if not device: + return "" + + try: + cmd = ["pm", "list", "packages", "-f"] + raw_output = str(device.shell(" ".join(cmd))) + except Exception as e: + logger.error(f"Failed to query package list from device: {e}") + return "" + + lines = raw_output.strip().split("\n") + packages = [] + for line in lines: + if "=" in line: + package_name = line.split("=")[-1].strip() + if package_name: + packages.append(package_name) + elif line.startswith("package:"): + package_name = line.split("package:")[-1].strip() + if package_name: + packages.append(package_name) + + return "\n".join(sorted(packages)) + + +async def list_packages_async(ctx: ArtemisContext) -> str: + return list_packages(ctx) + + +def get_current_foreground_package(ctx: ArtemisContext) -> str | None: + """Get the package name of the currently focused/foreground app dynamically.""" + device = get_adb_device(ctx) + if not device: + return None + + try: + app_info = device.current_app() + if app_info and app_info.package: + return app_info.package + except Exception as e: + logger.debug(f"device.current_app failed: {e}. Falling back to dumpsys.") + + try: + output = str(device.shell("dumpsys window | grep mCurrentFocus")) + if "mCurrentFocus=" in output: + segment = output.split("mCurrentFocus=")[-1] + if "/" in segment: + tokens = segment.split() + for token in tokens: + if "." in token and not token.startswith("Window"): + package = token.split("/")[0] + package = package.rstrip("}") + if package and "." in package: + return package + except Exception as e: + logger.debug(f"Fallback dumpsys parsing failed: {e}") + + return None + + +async def get_current_foreground_package_async(ctx: ArtemisContext) -> str | None: + return get_current_foreground_package(ctx) diff --git a/artemis/controllers/types.py b/artemis/controllers/types.py new file mode 100644 index 0000000..8d94c02 --- /dev/null +++ b/artemis/controllers/types.py @@ -0,0 +1,37 @@ +# 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. + +"""Backward compatibility alias for artemis.drivers.types.""" + +from artemis.drivers.types import ( + Bounds, + CoordinatesSelectorRequest, + PercentagesSelectorRequest, + SwipeRequest, + SwipeStartEndPercentagesRequest, + TapOutput, +) + +# Aliases for backward compatibility +SwipeStartEndCoordinatesRequest = SwipeRequest + +__all__ = [ + "Bounds", + "CoordinatesSelectorRequest", + "PercentagesSelectorRequest", + "SwipeRequest", + "SwipeStartEndCoordinatesRequest", + "SwipeStartEndPercentagesRequest", + "TapOutput", +] diff --git a/artemis/controllers/unified_controller.py b/artemis/controllers/unified_controller.py new file mode 100644 index 0000000..37afcbf --- /dev/null +++ b/artemis/controllers/unified_controller.py @@ -0,0 +1,796 @@ +# 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 asyncio +import os +from pathlib import Path +import tempfile +import time +from typing import Any +from uuid import uuid4 + +from artemis.config.paths import get_temp_dir +from artemis.context import ArtemisContext +from artemis.drivers.factory import get_driver +from artemis.drivers.base import BaseDeviceDriver +from artemis.controllers.device_controller import ScreenDataResponse +from artemis.controllers.types import ( + SwipeRequest, + SwipeStartEndCoordinatesRequest, + SwipeStartEndPercentagesRequest, + TapOutput, +) +from artemis.utils.logger import get_logger +from artemis.utils.video import ( + DEFAULT_MAX_DURATION_SECONDS, + RecordingSession, + VideoRecordingResult, + cleanup_video_segments, + concatenate_videos, + get_active_session, + get_ffmpeg_path, + has_active_session, + remove_active_session, + set_active_session, + trim_video, +) + +logger = get_logger(__name__) + + +class UnifiedMobileController: + def __init__(self, ctx: ArtemisContext): + self.ctx = ctx + self._driver: BaseDeviceDriver = get_driver(ctx) + self._segment_cache: dict[tuple[float, float | None], VideoRecordingResult] = {} + + @property + def driver(self) -> BaseDeviceDriver: + return self._driver + + @property + def controller(self) -> Any: + return self._driver + + async def tap_at( + self, + x: int, + y: int, + long_press: bool = False, + long_press_duration: int = 1000, + times: int = 1, + delay_ms: int = 100, + ) -> TapOutput: + try: + if long_press: + success = await self._driver.long_press(x, y, duration_ms=long_press_duration) + else: + success = await self._driver.tap( + x, y, duration_ms=100, times=times, delay_ms=delay_ms + ) + return TapOutput(error=None if success else f"Tap failed at ({x}, {y})") + except Exception as e: + return TapOutput(error=str(e)) + + async def tap_percentage( + self, + x_percent: int, + y_percent: int, + long_press: bool = False, + long_press_duration: int = 1000, + ) -> TapOutput: + """Tap at percentage-based coordinates (0 to 100).""" + norm_x = int(x_percent * 10) + norm_y = int(y_percent * 10) + success = await self._driver.tap_normalized( + norm_x, norm_y, long_press=long_press, duration_ms=long_press_duration + ) + return TapOutput( + error=None if success else f"Tap percentage failed at ({x_percent}%, {y_percent}%)" + ) + + async def tap_element( + self, + resource_id: str | None = None, + text: str | None = None, + index: int = 0, + long_press: bool = False, + long_press_duration: int = 1000, + ) -> TapOutput: + """Tap on a UI element by finding it in the hierarchy.""" + success = await self._driver.tap_element( + resource_id=resource_id, + text=text, + index=index, + long_press=long_press, + duration_ms=long_press_duration, + ) + return TapOutput( + error=None + if success + else f"Failed to tap element (resource_id={resource_id}, text={text})" + ) + + async def swipe_coords( + self, + start_x: int, + start_y: int, + end_x: int, + end_y: int, + duration: int = 400, + ) -> str | None: + """Swipe between two coordinate points.""" + success = await self._driver.swipe(start_x, start_y, end_x, end_y, duration_ms=duration) + return None if success else f"Swipe failed from ({start_x},{start_y}) to ({end_x},{end_y})" + + async def swipe_percentage( + self, + start_x_percent: int, + start_y_percent: int, + end_x_percent: int, + end_y_percent: int, + duration: int = 400, + ) -> str | None: + """Swipe using percentage-based coordinates (0 to 100).""" + start_norm = [int(start_x_percent * 10), int(start_y_percent * 10)] + end_norm = [int(end_x_percent * 10), int(end_y_percent * 10)] + success = await self._driver.swipe_normalized(start_norm, end_norm, duration_ms=duration) + return None if success else "Swipe percentage failed" + + async def swipe_request(self, request: SwipeRequest) -> str | None: + mode = request.swipe_mode + + if isinstance(mode, SwipeStartEndCoordinatesRequest): + return await self.swipe_coords( + start_x=mode.start.x, + start_y=mode.start.y, + end_x=mode.end.x, + end_y=mode.end.y, + duration=request.duration or 400, + ) + elif isinstance(mode, SwipeStartEndPercentagesRequest): + return await self.swipe_percentage( + start_x_percent=mode.start.x_percent, + start_y_percent=mode.start.y_percent, + end_x_percent=mode.end.x_percent, + end_y_percent=mode.end.y_percent, + duration=request.duration or 400, + ) + else: + return "Unsupported swipe mode" + + async def type_text(self, text: str, clear_existing: bool = True) -> bool: + return await self._driver.input_text(text, clear_existing=clear_existing) + + async def take_screenshot(self) -> str: + screen_data = await self._driver.get_screen_data() + return screen_data.screenshot_base64 + + async def launch_app(self, package_or_bundle_id: str) -> bool: + return await self._driver.launch_app(package_or_bundle_id) + + async def terminate_app(self, package_or_bundle_id: str | None) -> bool: + if not package_or_bundle_id: + return False + return await self._driver.stop_app(package_or_bundle_id) + + async def open_url(self, url: str) -> bool: + await self._driver.execute_shell(f"am start -a android.intent.action.VIEW -d '{url}'") + return True + + async def go_back(self) -> bool: + return await self._driver.press_key("back") + + async def go_home(self) -> bool: + return await self._driver.press_key("home") + + async def press_enter(self) -> bool: + return await self._driver.press_key("enter") + + async def press_key(self, keycode: str) -> bool: + return await self._driver.press_key(keycode) + + async def erase_text(self, nb_chars: int | None = None) -> bool: + if nb_chars is not None and nb_chars > 0: + for _ in range(nb_chars): + await self._driver.press_key("delete") + return True + # Safe & robust full clear: Move to End -> Shift+Home selection -> Delete -> Fallback backspaces + try: + clear_cmd = ( + "input keyevent 123 && " + "input keyevent --meta 1 122 && " + "input keyevent 67 && " + "input keyevent 67 67 67 67 67 67 67 67 67 67 67 67 67 67 67 67 67 67 67 67" + ) + await self._driver.execute_shell(clear_cmd) + except Exception: + for _ in range(30): + await self._driver.press_key("delete") + return True + + async def get_ui_elements(self) -> list[dict]: + screen_data = await self._driver.get_screen_data() + return screen_data.ui_elements or [] + + async def get_screen_data(self) -> "ScreenDataResponse": + """Get screen data including screenshot, UI hierarchy, dimensions, and platform.""" + data = await self._driver.get_screen_data() + return ScreenDataResponse( + base64=data.screenshot_base64, + elements=data.ui_elements or [], + width=data.width, + height=data.height, + platform=data.platform, + ) + + async def find_element( + self, + resource_id: str | None = None, + text: str | None = None, + index: int = 0, + ) -> tuple[dict | None, str | None]: + elem, _, error = await self._driver.find_element( + resource_id=resource_id, + text=text, + index=index, + ) + return elem, error + + def _get_device_id(self) -> str: + if self.ctx and self.ctx.device and self.ctx.device.device_id: + return self.ctx.device.device_id + return getattr(self._driver, "device_id", None) or "default" + + async def extract_segment_metadata( + self, + start_time: float, + end_time: float | None = None, + output_path: Path | None = None, + ) -> VideoRecordingResult: + """Get a video segment for a specific time range (relative to video start).""" + cache_key = ( + round(start_time, 1), + round(end_time, 1) if end_time is not None else None, + ) + if cache_key in self._segment_cache: + cached_res = self._segment_cache[cache_key] + if cached_res.success and cached_res.video_path and cached_res.video_path.exists(): + logger.info( + f"Reusing cached trimmed video segment for range {cache_key[0]}s to {cache_key[1]}s" + ) + return cached_res + + device_id = self._get_device_id() + + # Handle mock driver + if ( + getattr(self._driver, "is_mock", False) + or getattr(getattr(self.ctx, "device", None), "mobile_platform", None) == "mock" + or os.environ.get("ARTEMIS_MOCK_DRIVER") == "1" + ): + mock_video = Path("/tmp/mock_recording.mp4") + return VideoRecordingResult( + success=True, + video_path=mock_video, + message="Mock segment extracted", + ) + + session = get_active_session(device_id) + if not session: + return VideoRecordingResult( + success=False, + message=f"No active recording for device {device_id}", + ) + + try: + mkv_path = session.local_video_path + if not mkv_path: + return VideoRecordingResult( + success=False, + message="Recording file not found", + ) + + # Check if scrcpy process terminated and restart if needed (synchronous fallback) + crashed = False + restart_message = "" + if session.process and session.process.returncode is not None: + logger.warning(f"scrcpy process for {device_id} has terminated unexpectedly") + crashed = True + if mkv_path.exists() and mkv_path not in session.android_video_segments: + session.android_video_segments.append(mkv_path) + try: + session.android_segment_index += 1 + output_dir = mkv_path.parent + new_video_path = output_dir / f"recording_{session.android_segment_index}.mkv" + cmd = [ + "scrcpy", + "--serial", + device_id, + "--no-window", + "--record", + str(new_video_path), + "--record-format", + "mkv", + "--video-bit-rate", + "2M", + ] + process = await asyncio.create_subprocess_exec( + *cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + await asyncio.sleep(0.8) + if process.returncode is None: + session.process = process + session.local_video_path = new_video_path + mkv_path = new_video_path + restart_message = "Successfully restarted scrcpy." + else: + stderr = await process.stderr.read() + restart_message = f"Failed to restart scrcpy: {stderr.decode()}" + except Exception as e: + restart_message = f"Failed to restart scrcpy: {e}" + + # Determine source file to trim from (support multi-segment seamless trimming) + all_segments = [p for p in session.android_video_segments if p.exists()] + if mkv_path.exists() and mkv_path not in all_segments: + all_segments.append(mkv_path) + + if not all_segments: + return VideoRecordingResult( + success=False, + message="No recording segments found on disk", + ) + + temp_combined_file = None + if len(all_segments) > 1: + temp_combined_file = mkv_path.parent / "temp_combined_active.mkv" + concat_ok = await concatenate_videos(all_segments, temp_combined_file) + trim_source = ( + temp_combined_file if (concat_ok and temp_combined_file.exists()) else mkv_path + ) + else: + trim_source = all_segments[0] + + current_time = time.time() + video_duration = current_time - session.start_time + + # Strict timeline alignment: + # DataEngine session start time is the reference T0 (0.0s). + # Video recording start time is session.start_time. + offset = 0.0 + if session.data_engine_start_time is not None: + offset = max(0.0, session.start_time - session.data_engine_start_time) + + video_start_relative_time = start_time - offset + video_end_relative_time = end_time - offset if end_time is not None else None + + safe_duration = max(0.0, video_duration - 0.5) + + truncation_warning = None + if video_end_relative_time is not None and video_end_relative_time > safe_duration: + truncation_warning = ( + f"Video segment truncated. Requested end time {end_time:.1f}s " + f"exceeded available safe duration. Truncated to {safe_duration + offset:.1f}s." + ) + + if video_end_relative_time is None or video_end_relative_time > safe_duration: + video_end_relative_time = safe_duration + + if video_start_relative_time < 0: + video_start_relative_time = 0.0 + + if video_start_relative_time >= video_end_relative_time: + if safe_duration > 0: + video_end_relative_time = safe_duration + video_start_relative_time = max( + 0.0, + safe_duration - ((end_time - start_time) if end_time else 2.0), + ) + + if video_start_relative_time >= video_end_relative_time: + return VideoRecordingResult( + success=False, + message=( + f"Invalid time range or requested too close to current time. Available duration: {safe_duration:.1f}s" + ), + ) + + trim_output_dir = tempfile.mkdtemp( + prefix="video_trimmed_", + dir=get_temp_dir("trimmed_videos"), + ) + trim_output_path = Path(trim_output_dir) / "segment.mp4" + + success = False + for attempt in range(3): + success = await trim_video( + trim_source, + video_start_relative_time, + video_end_relative_time, + trim_output_path, + ) + if success and trim_output_path.exists(): + break + logger.warning(f"ffmpeg trim attempt {attempt + 1} failed, retrying in 0.5s...") + await asyncio.sleep(0.5) + + # Cleanup temp combined MKV if created + if temp_combined_file and temp_combined_file.exists(): + try: + temp_combined_file.unlink() + except Exception: + pass + + if not success or not trim_output_path.exists(): + return VideoRecordingResult( + success=False, + message="Failed to trim video segment after retries", + ) + + actual_start = video_start_relative_time + offset + actual_end = video_end_relative_time + offset + file_size_mb = trim_output_path.stat().st_size / (1024 * 1024) + duration = video_end_relative_time - video_start_relative_time + + message = f"Video segment retrieved for range {actual_start:.1f}s to {actual_end:.1f}s" + if crashed: + message += ( + f". Warning: scrcpy crashed. Data loss may have occurred. {restart_message}" + ) + + res = VideoRecordingResult( + success=True, + message=message, + video_path=trim_output_path, + file_size_mb=round(file_size_mb, 2), + duration_seconds=round(duration, 2), + actual_start_relative_time=actual_start, + warning=truncation_warning, + ) + self._segment_cache[cache_key] = res + return res + + except Exception as e: + logger.error(f"Failed to get video segment: {e}") + return VideoRecordingResult( + success=False, + message=f"Failed to get video segment: {e}", + ) + + async def _recording_watchdog(self, device_id: str) -> None: + """Background watchdog to automatically detect crashes and recover scrcpy recording.""" + while True: + session = get_active_session(device_id) + if not session or not session.is_active: + break + proc = session.process + if proc is None: + break + try: + ret = await proc.wait() + session = get_active_session(device_id) + if not session or not session.is_active: + break + + logger.warning( + f"scrcpy recording process for {device_id} terminated unexpectedly (code {ret}). " + f"Auto-restarting recording segment {session.android_segment_index + 1}..." + ) + if session.local_video_path and session.local_video_path.exists(): + if session.local_video_path not in session.android_video_segments: + session.android_video_segments.append(session.local_video_path) + + session.android_segment_index += 1 + output_dir = session.local_video_path.parent + new_video_path = output_dir / f"recording_{session.android_segment_index}.mkv" + cmd = [ + "scrcpy", + "--serial", + device_id, + "--no-window", + "--record", + str(new_video_path), + "--record-format", + "mkv", + "--video-bit-rate", + "2M", + ] + new_proc = await asyncio.create_subprocess_exec( + *cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + session.process = new_proc + session.local_video_path = new_video_path + await asyncio.sleep(0.8) + if new_proc.returncode is not None: + stderr = await new_proc.stderr.read() + logger.error(f"Failed to auto-restart scrcpy recording: {stderr.decode()}") + break + except asyncio.CancelledError: + break + except Exception as e: + logger.error(f"Error in recording watchdog for {device_id}: {e}") + break + + async def start_video_recording( + self, + output_dir: Path | None = None, + max_duration_seconds: int = DEFAULT_MAX_DURATION_SECONDS, + ) -> VideoRecordingResult: + """Start screen recording on Android device using scrcpy.""" + self._segment_cache.clear() + device_id = self._get_device_id() + + # Check mock driver first + if ( + getattr(self._driver, "is_mock", False) + or getattr(getattr(self.ctx, "device", None), "mobile_platform", None) == "mock" + or os.environ.get("ARTEMIS_MOCK_DRIVER") == "1" + ): + await self._driver.start_video_recording(output_dir) + return VideoRecordingResult(success=True, message="Mock recording started") + + if has_active_session(device_id): + return VideoRecordingResult( + success=False, + message=f"Recording already in progress for device {device_id}", + ) + + try: + if not output_dir: + output_dir = Path( + tempfile.mkdtemp(prefix="scrcpy_", dir=get_temp_dir("recordings")) + ) + else: + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + local_video_path = output_dir / "recording.mkv" + video_id = uuid4() + start_time = time.time() + data_engine_start_time = ( + self.ctx.data_engine.session_start_time + if (self.ctx and self.ctx.data_engine) + else start_time + ) + + session = RecordingSession( + video_id=video_id, + device_id=device_id, + start_time=start_time, + data_engine_start_time=data_engine_start_time, + local_video_path=local_video_path, + is_active=True, + ) + + # Persist to local database if Data Engine is active + if self.ctx and self.ctx.data_engine: + self.ctx.data_engine.record_video_start( + video_id=video_id, + device_id=device_id, + local_video_path=local_video_path, + start_time=session.start_time, + ) + + # Start scrcpy in background + cmd = [ + "scrcpy", + "--serial", + device_id, + "--no-window", + "--record", + str(local_video_path), + "--record-format", + "mkv", + "--video-bit-rate", + "2M", + ] + + process = await asyncio.create_subprocess_exec( + *cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + + session.process = process + set_active_session(device_id, session) + + logger.info(f"Started scrcpy recording on {device_id}, saving to {local_video_path}") + + # Check if process failed immediately + await asyncio.sleep(0.8) + if process.returncode is not None: + stderr = await process.stderr.read() + err_msg = stderr.decode() + logger.error(f"scrcpy failed to start on {device_id}: {err_msg}") + remove_active_session(device_id) + return VideoRecordingResult( + success=False, + message=f"scrcpy failed to start: {err_msg}", + ) + + # Start background watchdog to auto-recover if scrcpy terminates + session.watchdog_task = asyncio.create_task(self._recording_watchdog(device_id)) + + return VideoRecordingResult( + success=True, + message=f"Recording started on {device_id}", + ) + + except Exception as e: + logger.error(f"Failed to start scrcpy recording: {e}") + remove_active_session(device_id) + return VideoRecordingResult( + success=False, + message=f"Failed to start recording: {e}", + ) + + async def stop_video_recording(self) -> VideoRecordingResult: + """Stop scrcpy recording and return the converted MP4 video file.""" + self._segment_cache.clear() + device_id = self._get_device_id() + + # Check mock driver first + if ( + getattr(self._driver, "is_mock", False) + or getattr(getattr(self.ctx, "device", None), "mobile_platform", None) == "mock" + or os.environ.get("ARTEMIS_MOCK_DRIVER") == "1" + ): + p = await self._driver.stop_video_recording() + return VideoRecordingResult( + success=True, + video_path=Path(p) if p else None, + message="Mock recording stopped", + ) + + session = get_active_session(device_id) + if not session: + return VideoRecordingResult( + success=False, + message=f"No active recording for device {device_id}", + ) + + # Mark inactive and cancel watchdog task + session.is_active = False + if session.watchdog_task and not session.watchdog_task.done(): + session.watchdog_task.cancel() + + try: + process = session.process + if process is not None: + try: + process.terminate() + await asyncio.wait_for(process.wait(), timeout=5.0) + except Exception as proc_e: + logger.warning(f"Error terminating scrcpy process: {proc_e}") + + output_path = session.local_video_path + if not output_path or not output_path.exists(): + remove_active_session(device_id) + return VideoRecordingResult( + success=False, + message="Recording file not found on disk", + ) + + # Concatenate previous segments if scrcpy crashed and restarted + if session.android_video_segments: + all_segments = session.android_video_segments + if output_path.exists() and output_path not in all_segments: + all_segments.append(output_path) + concatenated_mkv = output_path.parent / "full_recording.mkv" + if await concatenate_videos(all_segments, concatenated_mkv): + cleanup_video_segments(all_segments) + output_path = concatenated_mkv + + # Convert MKV to web-streamable MP4 + mp4_path = ( + output_path.parent / "recording.mp4" + if output_path.name == "full_recording.mkv" + else output_path.with_suffix(".mp4") + ) + success = await self._convert_mkv_to_mp4(output_path, mp4_path) + + remove_active_session(device_id) + + final_video_path = mp4_path if (success and mp4_path.exists()) else output_path + + # Persist update to local database if Data Engine is active + if self.ctx and self.ctx.data_engine: + self.ctx.data_engine.record_video_stop( + video_id=session.video_id, + device_id=device_id, + local_video_path=final_video_path, + start_time=session.start_time, + end_time=time.time(), + ) + + if success and mp4_path.exists(): + try: + if output_path != mp4_path and output_path.exists(): + output_path.unlink() + except Exception: + pass + return VideoRecordingResult( + success=True, + message=f"Recording stopped, saved to {mp4_path}", + video_path=mp4_path, + ) + else: + return VideoRecordingResult( + success=True, + message=f"Recording stopped, saved as MKV to {output_path}", + video_path=output_path, + ) + + except Exception as e: + logger.error(f"Failed to stop scrcpy recording: {e}") + remove_active_session(device_id) + return VideoRecordingResult( + success=False, + message=f"Failed to stop recording: {e}", + ) + + async def _convert_mkv_to_mp4(self, mkv_path: Path, mp4_path: Path) -> bool: + """Convert MKV to MP4 using ffmpeg with faststart for web streaming.""" + if not mkv_path.exists(): + return False + try: + # 1. Try fast stream copy with +faststart + process = await asyncio.create_subprocess_exec( + get_ffmpeg_path(), + "-y", + "-i", + str(mkv_path), + "-c", + "copy", + "-movflags", + "+faststart", + str(mp4_path), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + await process.wait() + if mp4_path.exists() and mp4_path.stat().st_size > 0: + return True + + # 2. Fallback to re-encoding if stream copy fails + logger.warning("ffmpeg fast copy failed, re-encoding MKV to MP4...") + process = await asyncio.create_subprocess_exec( + get_ffmpeg_path(), + "-y", + "-i", + str(mkv_path), + "-c:v", + "libx264", + "-preset", + "fast", + "-pix_fmt", + "yuv420p", + "-movflags", + "+faststart", + str(mp4_path), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + await process.wait() + return mp4_path.exists() and mp4_path.stat().st_size > 0 + except Exception as e: + logger.error(f"Failed to convert MKV to MP4: {e}") + return False + + async def cleanup(self) -> None: + await self._driver.disconnect() diff --git a/artemis/core/__init__.py b/artemis/core/__init__.py new file mode 100644 index 0000000..c0dc5fc --- /dev/null +++ b/artemis/core/__init__.py @@ -0,0 +1,51 @@ +# 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. + +"""Core framework contracts and lifecycle primitives.""" + +from artemis.core.context import ExecutionContext, DeviceInfoContext +from artemis.core.state import ExecutionContextState, ExecutionStatus, StepRecord +from artemis.core.events import Event, EventBus, HookType, global_event_bus +from artemis.core.exceptions import ( + ArtemisException, + DeviceDriverException, + DeviceNotFoundException, + AgentExecutionException, + ToolExecutionException, + LLMProviderException, + StateTransitionException, +) +from artemis.core.registry import AgentRegistry, ToolRegistry, DriverRegistry + +__all__ = [ + "ExecutionContext", + "DeviceInfoContext", + "ExecutionContextState", + "ExecutionStatus", + "StepRecord", + "Event", + "EventBus", + "HookType", + "global_event_bus", + "ArtemisException", + "DeviceDriverException", + "DeviceNotFoundException", + "AgentExecutionException", + "ToolExecutionException", + "LLMProviderException", + "StateTransitionException", + "AgentRegistry", + "ToolRegistry", + "DriverRegistry", +] diff --git a/artemis/core/context.py b/artemis/core/context.py new file mode 100644 index 0000000..51a1b6d --- /dev/null +++ b/artemis/core/context.py @@ -0,0 +1,59 @@ +# 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. + +"""Clean Execution Context for ARTEMIS.""" + +from pathlib import Path +from typing import Any +from uuid import uuid4 +from pydantic import BaseModel, Field + + +class DeviceInfoContext(BaseModel): + """Device metadata and configuration.""" + + device_id: str = Field(default="default-device", description="ADB serial or device identifier") + platform: str = Field( + default="android", description="Platform type ('android', 'cloud', 'mock')" + ) + width: int = Field(default=1080, description="Screen width in pixels") + height: int = Field(default=2400, description="Screen height in pixels") + + +class ExecutionContext: + """Central context holding device driver, session identifiers, and telemetry tracer.""" + + def __init__( + self, + task_goal: str, + device_id: str = "default-device", + trace_id: str | None = None, + traces_path: Path | str | None = None, + locked_package: str | None = None, + platform: str = "android", + ): + self.task_goal = task_goal + self.trace_id = trace_id or str(uuid4()) + self.traces_path = Path(traces_path) if traces_path else Path("traces") / self.trace_id + self.locked_package = locked_package + self.device = DeviceInfoContext(device_id=device_id, platform=platform) + self.driver: Any | None = None + self.telemetry: Any | None = None + self.custom_attributes: dict[str, Any] = {} + + def set_attribute(self, key: str, value: Any) -> None: + self.custom_attributes[key] = value + + def get_attribute(self, key: str, default: Any = None) -> Any: + return self.custom_attributes.get(key, default) diff --git a/artemis/core/events.py b/artemis/core/events.py new file mode 100644 index 0000000..dd772b5 --- /dev/null +++ b/artemis/core/events.py @@ -0,0 +1,77 @@ +# 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. + +"""Event Bus and Lifecycle Hooks System for ARTEMIS.""" + +from collections import defaultdict +from collections.abc import Callable +from enum import Enum +import inspect +import time +from typing import Any +from pydantic import BaseModel, Field +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +class HookType(str, Enum): + BEFORE_TASK = "before_task" + AFTER_TASK = "after_task" + BEFORE_STEP = "before_step" + AFTER_STEP = "after_step" + ON_ACTION = "on_action" + ON_ERROR = "on_error" + ON_SCREEN_CHANGE = "on_screen_change" + + +class Event(BaseModel): + """Encapsulates an event payload emitted through the event bus.""" + + event_type: HookType = Field(..., description="Type of the lifecycle event") + timestamp: float = Field(default_factory=time.time, description="Event occurrence timestamp") + trace_id: str = Field(..., description="Active session trace identifier") + payload: dict[str, Any] = Field(default_factory=dict, description="Event data dictionary") + + +class EventBus: + """Asynchronous event bus supporting subscription and broadcast of lifecycle events.""" + + def __init__(self): + self._subscribers: dict[HookType, list[Callable[..., Any]]] = defaultdict(list) + + def subscribe(self, hook: HookType, callback: Callable[..., Any]) -> None: + """Registers a listener callback for a specific lifecycle event.""" + self._subscribers[hook].append(callback) + + def unsubscribe(self, hook: HookType, callback: Callable[..., Any]) -> None: + """Removes a registered callback.""" + if callback in self._subscribers[hook]: + self._subscribers[hook].remove(callback) + + async def emit(self, event: Event) -> None: + """Asynchronously dispatches an event to all subscribed listeners.""" + listeners = self._subscribers.get(event.event_type, []) + for listener in listeners: + try: + if inspect.iscoroutinefunction(listener): + await listener(event) + else: + listener(event) + except Exception as e: + logger.warning(f"Error executing event listener for {event.event_type}: {e}") + + +# Global singleton event bus instance +global_event_bus = EventBus() diff --git a/artemis/core/exceptions.py b/artemis/core/exceptions.py new file mode 100644 index 0000000..7eb1983 --- /dev/null +++ b/artemis/core/exceptions.py @@ -0,0 +1,57 @@ +# 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. + +"""Standardized exception hierarchy for ARTEMIS.""" + + +class ArtemisException(Exception): + """Base exception for all ARTEMIS framework errors.""" + + pass + + +class DeviceDriverException(ArtemisException): + """Raised when hardware, emulator, or ADB interactions fail.""" + + pass + + +class DeviceNotFoundException(DeviceDriverException): + """Raised when the specified target device cannot be located or connected.""" + + pass + + +class AgentExecutionException(ArtemisException): + """Raised when an autonomous agent fails during reasoning or execution.""" + + pass + + +class ToolExecutionException(ArtemisException): + """Raised when a tool action fails or receives invalid arguments.""" + + pass + + +class LLMProviderException(ArtemisException): + """Raised when an upstream LLM/VLM API call fails or times out.""" + + pass + + +class StateTransitionException(ArtemisException): + """Raised when graph execution encounters an invalid state transition.""" + + pass diff --git a/artemis/core/registry.py b/artemis/core/registry.py new file mode 100644 index 0000000..140e877 --- /dev/null +++ b/artemis/core/registry.py @@ -0,0 +1,46 @@ +# 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. + +"""Global Plugin and Component Registry for Agents, Tools, and Drivers.""" + +from typing import Any, TypeVar + +T = TypeVar("T") + + +class ComponentRegistry: + """Generic typed component registry supporting runtime dynamic plugin loading.""" + + def __init__(self, name: str): + self.name = name + self._items: dict[str, Any] = {} + + def register(self, key: str, item: Any) -> Any: + self._items[key.lower()] = item + return item + + def get(self, key: str) -> Any | None: + return self._items.get(key.lower()) + + def list_all(self) -> dict[str, Any]: + return dict(self._items) + + def contains(self, key: str) -> bool: + return key.lower() in self._items + + +# Global Registries +AgentRegistry = ComponentRegistry("Agents") +ToolRegistry = ComponentRegistry("Tools") +DriverRegistry = ComponentRegistry("Drivers") diff --git a/artemis/core/state.py b/artemis/core/state.py new file mode 100644 index 0000000..9c3a4e0 --- /dev/null +++ b/artemis/core/state.py @@ -0,0 +1,70 @@ +# 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. + +"""Core Execution State and Lifecycle Data Structures.""" + +from enum import Enum +import time +from typing import Any +from pydantic import BaseModel, Field + + +class ExecutionStatus(str, Enum): + PENDING = "pending" + RUNNING = "running" + SUCCESS = "success" + FAILED = "failed" + PAUSED = "paused" + INTERRUPTED = "interrupted" + + +class StepRecord(BaseModel): + """Encapsulates a single step's observation, thought, and action.""" + + step_number: int = Field(..., description="1-indexed step turn counter") + thought: str | None = Field(default=None, description="Reasoning text") + action_name: str | None = Field(default=None, description="Triggered action name") + action_params: dict[str, Any] = Field(default_factory=dict, description="Action arguments") + result: str | None = Field(default=None, description="Action execution result output") + screenshot_path: str | None = Field(default=None, description="Saved step screenshot file URI") + timestamp: float = Field(default_factory=time.time, description="Timestamp of step start") + duration_seconds: float = Field(default=0.0, description="Step processing duration in seconds") + + +class ExecutionContextState(BaseModel): + """Full execution state passed across agents and runners.""" + + task_goal: str = Field(..., description="Initial user task objective") + trace_id: str = Field(..., description="Unique task trace identifier") + status: ExecutionStatus = Field( + default=ExecutionStatus.RUNNING, description="Current execution lifecycle status" + ) + current_turn: int = Field(default=0, description="Current step turn count") + max_turns: int = Field(default=30, description="Maximum allowed step turns before timeout") + steps: list[StepRecord] = Field( + default_factory=list, description="Chronological execution step history" + ) + active_plan: list[str] = Field( + default_factory=list, description="Decomposed milestone plan items" + ) + current_step_index: int = Field(default=0, description="Current index within active_plan") + metadata: dict[str, Any] = Field( + default_factory=dict, description="Custom metadata and state parameters" + ) + error_message: str | None = Field( + default=None, description="Failure reason if status is FAILED" + ) + final_output: Any | None = Field( + default=None, description="Final synthesized structured output from Outputter or runner" + ) diff --git a/artemis/core/tool_declaration.py b/artemis/core/tool_declaration.py new file mode 100644 index 0000000..c40acae --- /dev/null +++ b/artemis/core/tool_declaration.py @@ -0,0 +1,46 @@ +# 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. + +"""Standardized tool declaration adhering to JSON Schema format.""" + +from typing import Any + + +class ToolDeclaration(dict): + """Standardized tool declaration adhering to JSON Schema format with property accessors.""" + + def __init__(self, name: str, description: str, parameters: dict[str, Any]): + super().__init__( + type="function", + function={ + "name": name, + "description": description, + "parameters": parameters, + }, + name=name, + description=description, + parameters=parameters, + ) + + @property + def name(self) -> str: + return self.get("name") or self.get("function", {}).get("name", "") + + @property + def description(self) -> str: + return self.get("description") or self.get("function", {}).get("description", "") + + @property + def parameters(self) -> dict[str, Any]: + return self.get("parameters") or self.get("function", {}).get("parameters", {}) diff --git a/artemis/data_engine/context_vars.py b/artemis/data_engine/context_vars.py new file mode 100644 index 0000000..4a39f4b --- /dev/null +++ b/artemis/data_engine/context_vars.py @@ -0,0 +1,24 @@ +# 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. + +"""Context variables for Artemis DataEngine runtime tracing.""" + +import contextvars + +CURRENT_TRACE_ID: contextvars.ContextVar[str | None] = contextvars.ContextVar( + "current_trace_id", default=None +) +CURRENT_NODE_NAME: contextvars.ContextVar[str | None] = contextvars.ContextVar( + "current_node_name", default=None +) diff --git a/artemis/data_engine/engine.py b/artemis/data_engine/engine.py new file mode 100644 index 0000000..f7c7c12 --- /dev/null +++ b/artemis/data_engine/engine.py @@ -0,0 +1,1138 @@ +# 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 asyncio +from collections.abc import Callable +import hashlib +import json +import os +from pathlib import Path +import socket +import threading +import time +from typing import Any +from uuid import UUID, uuid4 + +from artemis.config import read_ipc_port, settings +from artemis.context import ArtemisContext +from artemis.data_engine.models import ( + BackgroundTaskRecord, + FailedOutputRecord, + ImageRecord, + SessionMetadata, + StepRecord, + TraceRecord, + VideoRecordingRecord, +) +from artemis.data_engine.storage import StorageManager +from artemis.data_engine.trace import CURRENT_TRACE_ID +from artemis.utils.coordinates import ( + normalize_any_structure, + normalize_step_actions, +) +from artemis.utils.logger import get_logger +from artemis.utils.text import safe_extract_text + +logger = get_logger(__name__) + +_CURRENT_DATA_ENGINE = None + + +class DataEngine: + """Core engine for managing runtime data, storage, and streaming.""" + + def __init__(self, ctx: ArtemisContext): + self.ctx = ctx + + # Determine paths from context + if ctx.execution_setup and ctx.execution_setup.traces_path: + self.global_base_dir = Path(ctx.execution_setup.traces_path) + self.db_path = self.global_base_dir / "data_engine.db" + else: + # Fallback to default unified settings + self.global_base_dir = settings.TRACES_PATH + self.db_path = settings.DATA_ENGINE_DB_PATH + + self.storage = StorageManager(self.db_path, self.global_base_dir) + + self.current_session_id: UUID | None = None + self.session_start_time: float | None = None + self.current_step_id: UUID | None = None + self.last_recorded_step_id: UUID | None = None + self.current_step_dir: Path | None = None + self.subscribers: list[Callable[[str, Any], None]] = [] + self.lock = threading.Lock() + self._trace_counter = 0 + self._pending_tasks = set() + self._pending_threads = [] + self._accumulated_logs = {} + self._bg_task_to_trace_id = {} + # Background tasks are written directly to SQLite storage + + self.ipc_socket = None + + ipc_port = read_ipc_port() + if ipc_port: + try: + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + s.connect(("127.0.0.1", ipc_port)) + s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + self.ipc_socket = s + logger.info(f"Connected to IPC server on port {ipc_port} with TCP_NODELAY") + except Exception as e: + logger.error(f"Failed to connect to IPC server on port {ipc_port}: {e}") + + @property + def base_dir(self) -> Path: + if self.current_session_id: + return self.global_base_dir / str(self.current_session_id) + return self.global_base_dir + + def subscribe(self, callback: Callable[[str, Any], None]): + """Subscribe to real-time events (e.g., for SSE).""" + self.subscribers.append(callback) + + def unsubscribe(self, callback: Callable[[str, Any], None]): + """Unsubscribe from real-time events.""" + if callback in self.subscribers: + self.subscribers.remove(callback) + + def _publish(self, event_type: str, data: Any): + """Publish event to all subscribers.""" + if isinstance(data, dict) and "session_id" not in data and self.current_session_id: + data["session_id"] = str(self.current_session_id) + + for callback in self.subscribers: + try: + callback(event_type, data) + except Exception as e: + logger.error(f"Error in subscriber callback: {e}") + + # Bridge to Cloud Gateway SSE stream in Cloud Mode + if os.environ.get("ARTEMIS_CLOUD_MODE") == "1": + try: + import asyncio + from debug_pub.cloud_gateway import cloud_manager + + sid = os.environ.get("ARTEMIS_CLOUD_SESSION_ID") or ( + str(self.current_session_id) if self.current_session_id else None + ) + if sid and sid in cloud_manager.sessions: + try: + loop = asyncio.get_running_loop() + if loop.is_running(): + loop.create_task( + cloud_manager.sessions[sid].emit_event( + event_type, + data if isinstance(data, dict) else {"payload": data}, + ) + ) + except RuntimeError: + pass + except Exception: + pass + + # Auto-reconnect to IPC socket if disconnected or not yet connected + if not getattr(self, "ipc_socket", None): + ipc_port = read_ipc_port() + if ipc_port: + try: + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + s.connect(("127.0.0.1", ipc_port)) + s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + self.ipc_socket = s + except Exception: + self.ipc_socket = None + + if getattr(self, "ipc_socket", None): + try: + payload = json.dumps({"event_type": event_type, "data": data}, default=str) + "\n" + self.ipc_socket.sendall(payload.encode("utf-8")) + except Exception: + try: + self.ipc_socket.close() + except Exception: + pass + self.ipc_socket = None + + def start_session(self, goal: str, device_info: dict[str, Any] | None = None) -> UUID: + """Start a new session.""" + global _CURRENT_DATA_ENGINE + _CURRENT_DATA_ENGINE = self + + env_session_id = os.getenv("ARTEMIS_CLOUD_SESSION_ID") or os.getenv("ARTEMIS_SESSION_ID") + if env_session_id: + try: + session_id = UUID(env_session_id) + except ValueError: + session_id = env_session_id + else: + session_id = uuid4() + self.current_session_id = session_id + self.session_start_time = time.time() + + # Clear old pause file if it exists + pause_file = self.global_base_dir.parent / ".artemis_paused" + if pause_file.exists(): + try: + pause_file.unlink() + logger.info("Removed old pause file on session start.") + except Exception as e: + logger.error(f"Failed to delete old pause file: {e}") + + session = SessionMetadata( + session_id=session_id, + initial_goal=goal, + start_time=self.session_start_time, + device_info=device_info or {}, + pid=os.getpid(), + ) + self.storage.create_session(session) + self.current_step_number = 0 + logger.info(f"Session started: {session_id}") + self._publish("session_started", session.model_dump()) + return session_id + + def end_session(self, status: str = "completed"): + """End the current session, updating its status and end time.""" + if not self.current_session_id: + return + + # Clear pause file if it exists + if hasattr(self, "global_base_dir") and self.global_base_dir: + pause_file = self.global_base_dir.parent / ".artemis_paused" + if pause_file.exists(): + try: + pause_file.unlink() + logger.info("Removed pause file on session end.") + except Exception as e: + logger.error(f"Failed to delete pause file on session end: {e}") + + session_id = self.current_session_id + end_time = time.time() + session = self.storage.get_session(session_id) + if session: + session.end_time = end_time + session.status = status + else: + session = SessionMetadata( + session_id=session_id, + initial_goal="", + start_time=self.session_start_time or end_time, + end_time=end_time, + status=status, + device_info=self.ctx.device.model_dump() + if getattr(self, "ctx", None) and self.ctx.device + else {}, + ) + + try: + self.storage.update_session(session) + logger.info(f"Session ended: {session_id} with status: {status}") + self._publish("session_ended", session.model_dump()) + except Exception as e: + logger.error(f"Failed to end session in DataEngine: {e}") + + def record_video_start( + self, + video_id: UUID, + device_id: str, + local_video_path: str | Path, + start_time: float | None = None, + ): + """Record the start of a video recording.""" + if not self.storage: + return + try: + record = VideoRecordingRecord( + video_id=video_id, + session_id=self.current_session_id, + device_id=device_id, + start_time=start_time or time.time(), + local_video_path=str(local_video_path), + ) + self.storage.create_video_recording(record) + except Exception as e: + logger.error(f"Failed to record video start in DataEngine: {e}") + + def record_video_stop( + self, + video_id: UUID, + device_id: str, + local_video_path: str | Path, + start_time: float, + end_time: float | None = None, + ): + """Record the completion of a video recording and sync to session metadata.""" + if not self.storage: + return + try: + record = VideoRecordingRecord( + video_id=video_id, + session_id=self.current_session_id, + device_id=device_id, + start_time=start_time, + end_time=end_time or time.time(), + local_video_path=str(local_video_path), + ) + self.storage.update_video_recording(record) + except Exception as e: + logger.error(f"Failed to record video stop in DataEngine: {e}") + + def update_video_path(self, local_video_path: str | Path): + """Update the video path across tables when traces or videos are moved.""" + if not self.storage or not self.current_session_id: + return + try: + self.storage.update_session_video_path(self.current_session_id, str(local_video_path)) + except Exception as e: + logger.warning(f"Failed to update video path in DataEngine: {e}") + + def get_or_create_image( + self, + image_bytes: bytes, + ui_tree: Any | None = None, + ocr_result: Any | None = None, + ) -> str: + """Get image name by hash, or create new record if not exists.""" + + hasher = hashlib.sha256() + hasher.update(image_bytes) + image_name = hasher.hexdigest() + + image_record = self.storage.get_image(image_name) + if image_record: + # If the image exists but is missing OCR or UI tree, and we now have them, update the record. + if (ocr_result is not None and image_record.ocr_result is None) or ( + ui_tree is not None and image_record.ui_tree is None + ): + self.storage.update_image_data(image_name, ocr_result, ui_tree) + return image_name + + images_dir = self.global_base_dir / "images" + images_dir.mkdir(parents=True, exist_ok=True) + file_path = images_dir / f"{image_name}.jpg" + + with open(file_path, "wb") as f: + f.write(image_bytes) + + new_record = ImageRecord( + image_name=image_name, + ui_tree=ui_tree, + ocr_result=ocr_result, + extra_metadata={}, + ) + self.storage.create_image(new_record) + + return image_name + + def get_image_path(self, image_name: str) -> Path: + """Get the absolute file path of an image by its name/hash.""" + return self.global_base_dir / "images" / f"{image_name}.jpg" + + def record_step( + self, + pre_screenshot_bytes: bytes | None = None, + post_screenshot_bytes: bytes | None = None, + ui_tree: Any | None = None, + ocr_result: Any | None = None, + foreground_app: str | None = None, + action_taken: dict[str, Any] | None = None, + summary: str | None = None, + operator_raw_thinking: str | None = None, + operator_native_thinking: str | None = None, + last_execution_result: Any | None = None, + extra_metadata: dict[str, Any] | None = None, + ) -> UUID: + """Record a step, saving screenshots and metadata.""" + if not self.current_session_id: + raise ValueError("No active session. Call start_session first.") + + with self.lock: + self.current_step_number += 1 + step_number = self.current_step_number + + step_id = self.current_step_id or uuid4() + self.last_recorded_step_id = step_id + # Reset current_step_id so subsequent steps will allocate a new one + self.current_step_id = None + + # Process images and step creation in background thread + + pre_image_name = ( + hashlib.sha256(pre_screenshot_bytes).hexdigest() if pre_screenshot_bytes else None + ) + post_image_name = ( + hashlib.sha256(post_screenshot_bytes).hexdigest() if post_screenshot_bytes else None + ) + if pre_image_name and post_image_name and pre_image_name == post_image_name: + post_image_name = None + + step = StepRecord( + step_id=step_id, + session_id=self.current_session_id, + step_number=step_number, + timestamp=time.time(), + pre_image_name=pre_image_name, + post_image_name=post_image_name, + summary=summary, + action_taken=action_taken, + operator_raw_thinking=operator_raw_thinking, + operator_native_thinking=operator_native_thinking, + last_execution_result=last_execution_result, + extra_metadata=extra_metadata or {}, + ) + + def _run_background_storage(): + if pre_screenshot_bytes: + self.get_or_create_image(pre_screenshot_bytes, ui_tree, ocr_result) + if post_screenshot_bytes: + self.get_or_create_image(post_screenshot_bytes) + self.storage.create_step(step) + + # Run storage in background + try: + task = asyncio.create_task(asyncio.to_thread(_run_background_storage)) + self._pending_tasks.add(task) + task.add_done_callback(self._on_background_task_done) + except RuntimeError: + # Fallback if no event loop + def run_and_cleanup(): + _run_background_storage() + curr_thread = threading.current_thread() + if curr_thread in self._pending_threads: + self._pending_threads.remove(curr_thread) + + thread = threading.Thread(target=run_and_cleanup) + self._pending_threads.append(thread) + thread.start() + + logger.info(f"Recorded step {step_number} for session {self.current_session_id}") + + # Publish event with relative time + step_dict = step.model_dump() + step_dict["relative_time"] = self.get_relative_time(step.timestamp) + + # Include generic tools in the real-time event + try: + traces = self.storage.get_step_traces(step_id) + generic_tools = [] + for t in traces: + if t.type == "tool": + generic_tools.append(t.model_dump()) + step_dict["generic_tools"] = generic_tools + except Exception as e: + logger.error(f"Failed to fetch step traces for SSE: {e}") + step_dict["generic_tools"] = [] + + # Extract and attach token usage for real-time SSE stream + step_p, step_c, step_t = 0, 0, 0 + if extra_metadata and isinstance(extra_metadata, dict) and "token_usage" in extra_metadata: + u = extra_metadata["token_usage"] + if isinstance(u, dict): + step_p = int(u.get("prompt_tokens") or u.get("input_tokens") or 0) + step_c = int(u.get("completion_tokens") or u.get("output_tokens") or 0) + step_t = int(u.get("total_tokens") or (step_p + step_c)) + + if step_t == 0: + token_info = self._get_step_token_usage_for_sse(step_id) + if token_info: + step_p = token_info["prompt_tokens"] + step_c = token_info["completion_tokens"] + step_t = token_info["total_tokens"] + + if step_t > 0: + step_dict["token_usage"] = { + "prompt_tokens": step_p, + "completion_tokens": step_c, + "total_tokens": step_t, + } + step_dict["total_tokens"] = step_t + + step_dict = normalize_step_actions(step_dict) + + self._publish("step_recorded", step_dict) + + return step_id + + def record_trace( + self, + type: str, + name: str, + payload: dict[str, Any], + step_id: UUID | None = None, + parent_trace_id: UUID | None = None, + status: str = "success", + duration: float | None = None, + trace_id: UUID | None = None, + ) -> UUID: + """Record a trace (agent, tool, or log), non-blocking.""" + if not self.current_session_id: + raise ValueError("No active session. Call start_session first.") + + trace_id = trace_id or uuid4() + + with self.lock: + self._trace_counter = getattr(self, "_trace_counter", 0) + 1 + trace_ts = time.time() + (self._trace_counter * 1e-7) + + trace = TraceRecord( + trace_id=trace_id, + session_id=self.current_session_id, + step_id=step_id, + parent_trace_id=parent_trace_id, + type=type, + name=name, + timestamp=trace_ts, + duration=duration, + status=status, + payload=payload, + ) + + # Run storage in background + try: + asyncio.get_running_loop() + task = asyncio.create_task(asyncio.to_thread(self.storage.create_trace, trace)) + self._pending_tasks.add(task) + task.add_done_callback(self._on_background_task_done) + except RuntimeError: + # Fallback if no event loop + def run_and_cleanup(): + self.storage.create_trace(trace) + curr_thread = threading.current_thread() + if curr_thread in self._pending_threads: + self._pending_threads.remove(curr_thread) + + thread = threading.Thread(target=run_and_cleanup) + self._pending_threads.append(thread) + thread.start() + + # Publish event + self._publish("trace_recorded", trace.model_dump()) + + return trace_id + + def record_failed_output( + self, + trace_id: UUID, + model_name: str, + prompt: str, + raw_output: str, + error_message: str, + ): + """Record a malformed output (crime scene) for later training.""" + if not self.current_session_id: + return + + record = FailedOutputRecord( + session_id=self.current_session_id, + trace_id=trace_id, + model_name=model_name, + prompt=prompt, + raw_output=raw_output, + error_message=error_message, + ) + + try: + task = asyncio.create_task(asyncio.to_thread(self.storage.create_failed_output, record)) + self._pending_tasks.add(task) + task.add_done_callback(self._pending_tasks.discard) + except RuntimeError: + + def run_and_cleanup(): + self.storage.create_failed_output(record) + curr_thread = threading.current_thread() + if curr_thread in self._pending_threads: + self._pending_threads.remove(curr_thread) + + thread = threading.Thread(target=run_and_cleanup) + self._pending_threads.append(thread) + thread.start() + + def has_pending_operations(self) -> bool: + """Check if there are any pending background tasks or threads.""" + return ( + len(self._pending_tasks) > 0 + or len([t for t in self._pending_threads if t.is_alive()]) > 0 + ) + + async def shutdown(self): + """Wait for all pending background tasks and threads to complete.""" + if self._pending_tasks: + logger.info( + f"Waiting for {len(self._pending_tasks)} pending async tasks to complete..." + ) + await asyncio.gather(*self._pending_tasks, return_exceptions=True) + + if self._pending_threads: + logger.info(f"Waiting for {len(self._pending_threads)} pending threads to complete...") + threads_to_join = list(self._pending_threads) + for thread in threads_to_join: + if thread.is_alive(): + await asyncio.to_thread(thread.join) + + if getattr(self, "ipc_socket", None): + try: + self.ipc_socket.close() + except Exception: + pass + + logger.info("DataEngine shutdown complete. All data persisted.") + + def stream_output( + self, + execution_id: UUID, + chunk: str, + stream_type: str | None = None, + is_thinking: bool | None = None, + ): + """Stream LLM output chunks.""" + if stream_type is None: + if is_thinking is True: + stream_type = "thinking" + else: + stream_type = "text" + + parent_id = CURRENT_TRACE_ID.get() + session_id_str = str(self.current_session_id) if self.current_session_id else None + + self._publish( + "llm_stream", + { + "execution_id": str(execution_id), + "session_id": session_id_str, + "parent_trace_id": str(parent_id) if parent_id else None, + "step_id": (str(self.current_step_id) if self.current_step_id else None), + "chunk": chunk, + "stream_type": stream_type, + "is_thinking": ( + is_thinking if is_thinking is not None else (stream_type == "thinking") + ), + }, + ) + + exec_str = str(execution_id) + if not hasattr(self, "_accumulated_logs"): + self._accumulated_logs = {} + if exec_str not in self._accumulated_logs: + self._accumulated_logs[exec_str] = "" + self._accumulated_logs[exec_str] += chunk + + if session_id_str: + if session_id_str not in self._accumulated_logs: + self._accumulated_logs[session_id_str] = "" + self._accumulated_logs[session_id_str] += chunk + + def _get_generic_tools_for_sse(self, step_id: UUID) -> list[dict[str, Any]]: + try: + traces = self.storage.get_step_traces(step_id) + return [t.model_dump() for t in traces if t.type == "tool"] + except Exception as e: + logger.error(f"Failed to fetch step traces for SSE: {e}") + return [] + + def _get_step_token_usage_for_sse(self, step_id: UUID) -> dict[str, Any] | None: + try: + traces = self.storage.get_step_traces(step_id) + step_p, step_c, step_t = 0, 0, 0 + for t in traces: + if t.type == "llm_call" and isinstance(t.payload, dict): + u = t.payload.get("token_usage") or t.payload.get("usage_metadata") + if isinstance(u, dict): + p = int( + u.get("prompt_tokens") + or u.get("prompt_token_count") + or u.get("input_tokens") + or 0 + ) + c = int( + u.get("completion_tokens") + or u.get("candidates_token_count") + or u.get("output_tokens") + or 0 + ) + tot = int(u.get("total_tokens") or u.get("total_token_count") or (p + c)) + step_p += p + step_c += c + step_t += tot + if step_t > 0: + return { + "prompt_tokens": step_p, + "completion_tokens": step_c, + "total_tokens": step_t, + } + except Exception: + pass + return None + + def _on_background_task_done(self, task: asyncio.Task): + self._pending_tasks.discard(task) + if not task.cancelled() and task.exception(): + exc = task.exception() + logger.error(f"Background task failed with exception: {exc}", exc_info=exc) + raise exc + + def _run_in_background(self, fn, *args): + """Safely offload blocking storage/disk operation to background without stalling event loop.""" + try: + task = asyncio.create_task(asyncio.to_thread(fn, *args)) + self._pending_tasks.add(task) + task.add_done_callback(self._on_background_task_done) + except RuntimeError: + + def run_and_cleanup(): + try: + fn(*args) + finally: + curr_thread = threading.current_thread() + if curr_thread in self._pending_threads: + self._pending_threads.remove(curr_thread) + + thread = threading.Thread(target=run_and_cleanup) + self._pending_threads.append(thread) + thread.start() + + def update_step_action( + self, + action_taken: dict[str, Any], + post_screenshot_bytes: bytes | None = None, + ): + """Update the current step with action taken and post-action screenshot in background.""" + with self.lock: + step_id = self.current_step_id or self.last_recorded_step_id + if not step_id: + raise ValueError("No active step to update.") + if self.base_dir: + step_dir = self.base_dir / "steps" / str(step_id) + self.current_step_dir = step_dir + step_dir.mkdir(parents=True, exist_ok=True) + else: + step_dir = None + + def _update_and_write(): + self.storage.update_step_action(step_id, action_taken) + if post_screenshot_bytes and step_dir: + step_dir.mkdir(parents=True, exist_ok=True) + with open(step_dir / "post.jpg", "wb") as f: + f.write(post_screenshot_bytes) + + self._run_in_background(_update_and_write) + + update_payload = { + "step_id": str(step_id), + "action_taken": action_taken, + "generic_tools": self._get_generic_tools_for_sse(step_id), + } + tokens = self._get_step_token_usage_for_sse(step_id) + if tokens: + update_payload["token_usage"] = tokens + update_payload["total_tokens"] = tokens["total_tokens"] + + self._publish("step_updated", update_payload) + + def update_step_summary(self, step_id: UUID | int | str, summary: str): + """Update the step with a concise summary in background.""" + target_uuid = None + if isinstance(step_id, int): + if self.current_session_id: + steps = self.storage.get_steps(self.current_session_id) + for s in steps: + if s.step_number == step_id: + target_uuid = s.step_id + break + elif isinstance(step_id, str): + try: + target_uuid = UUID(step_id) + except ValueError: + target_uuid = step_id + else: + target_uuid = step_id + + if not target_uuid: + logger.debug(f"Could not resolve step_id for {step_id} to update summary.") + return + + self._run_in_background(self.storage.update_step_summary, target_uuid, summary) + self._publish( + "step_updated", + { + "step_id": str(target_uuid), + "summary": summary, + "generic_tools": self._get_generic_tools_for_sse(target_uuid), + }, + ) + + def update_step_thinking(self, step_id: UUID, operator_raw_thinking: str): + """Update step's thinking in SQLite in background.""" + self._run_in_background(self.storage.update_step_thinking, step_id, operator_raw_thinking) + self._publish( + "step_updated", + { + "step_id": str(step_id), + "operator_raw_thinking": operator_raw_thinking, + "generic_tools": self._get_generic_tools_for_sse(step_id), + }, + ) + + def update_step_native_thinking(self, step_id: UUID, operator_native_thinking: str): + """Update step's native thinking in SQLite in background.""" + self._run_in_background( + self.storage.update_step_native_thinking, + step_id, + operator_native_thinking, + ) + self._publish( + "step_updated", + { + "step_id": str(step_id), + "operator_native_thinking": operator_native_thinking, + "generic_tools": self._get_generic_tools_for_sse(step_id), + }, + ) + + def update_step_execution_result( + self, + step_id: UUID, + last_execution_result: dict, + post_image_name: str | None = None, + ): + """Update step's execution result and post_image_name in SQLite in background.""" + self._run_in_background( + self.storage.update_step_execution_result, + step_id, + last_execution_result, + post_image_name, + ) + self._publish( + "step_updated", + { + "step_id": str(step_id), + "last_execution_result": last_execution_result, + "post_image_name": post_image_name, + "generic_tools": self._get_generic_tools_for_sse(step_id), + }, + ) + + def get_relative_time(self, timestamp: float) -> str: + """Compute relative time since session start.""" + if self.session_start_time is None: + return "0.0s" + diff = timestamp - self.session_start_time + return f"{diff:.1f}s" + + def allocate_step_id(self) -> UUID: + """Pre-allocate a step ID for the upcoming step to synchronize traces.""" + self.current_step_id = uuid4() + if self.base_dir: + self.current_step_dir = self.base_dir / "steps" / str(self.current_step_id) + self.current_step_dir.mkdir(parents=True, exist_ok=True) + return self.current_step_id + + def get_agent_friendly_steps(self) -> list[dict[str, Any]]: + """Retrieve steps formatted for agent consumption (relative time).""" + if not self.current_session_id: + return [] + steps_with_traces = self.storage.get_steps_with_traces(self.current_session_id) + result = [] + for step, traces in steps_with_traces: + step_dict = step.model_dump() + step_dict["relative_time"] = self.get_relative_time(step.timestamp) + + interleaved_events = [] + try: + blacklist = ( + "operator", + "perception", + "planner", + "validator", + "summarizer", + "checker", + ) + + relevant_traces = [] + for t in traces: + if t.status not in ("success", "failed"): + continue + + is_relevant_llm = False + if t.type == "llm_call": + curr_parent = t.parent_trace_id + visited = set() + while curr_parent: + if curr_parent in visited: + break + visited.add(curr_parent) + parent_trace = next( + (x for x in traces if x.trace_id == curr_parent), + None, + ) + if parent_trace: + if parent_trace.name in ( + "operator", + "failure_analyzer", + ): + is_relevant_llm = True + t.name = parent_trace.name + break + curr_parent = parent_trace.parent_trace_id + else: + break + + if is_relevant_llm: + relevant_traces.append(t) + elif t.type in ("tool", "agent") and t.name not in blacklist: + relevant_traces.append(t) + + # Sort by timestamp to ensure exact chronological order + relevant_traces.sort(key=lambda x: x.timestamp) + + # Separate set of candidate tool/agent IDs to do top-level filtering + # (only keep top-level tool calls) + candidate_ids = {t.trace_id for t in relevant_traces if t.type in ("tool", "agent")} + + for t in relevant_traces: + if t.type == "llm_call": + payload = t.payload or {} + response_list = payload.get("response") or [] + thought_text = "" + has_structured_blocks = False + if response_list: + first_gen = response_list[0] + if isinstance(first_gen, dict): + content_val = first_gen.get("content") + if isinstance(content_val, list): + has_structured_blocks = True + for block in content_val: + if isinstance(block, dict): + if ( + block.get("type") == "thinking" + and block.get("thinking", "").strip() + ): + e_type = ( + "failure_analyzer_native_thought" + if t.name == "failure_analyzer" + else "native_thought" + ) + interleaved_events.append( + { + "type": e_type, + "content": (block["thinking"].strip()), + } + ) + elif ( + block.get("type") == "text" + and block.get("text", "").strip() + ): + e_type = ( + "failure_analyzer_thought" + if t.name == "failure_analyzer" + else "thought" + ) + interleaved_events.append( + { + "type": e_type, + "content": (block["text"].strip()), + } + ) + if not has_structured_blocks: + thought_text = ( + first_gen.get("content") or first_gen.get("text") or "" + ) + else: + thought_text = str(first_gen) + + if not has_structured_blocks: + if isinstance(thought_text, list): + thought_text = safe_extract_text(thought_text) + + if thought_text and thought_text.strip(): + e_type = ( + "failure_analyzer_thought" + if t.name == "failure_analyzer" + else "thought" + ) + interleaved_events.append( + { + "type": e_type, + "content": thought_text.strip(), + } + ) + elif t.type in ("tool", "agent"): + # Keep only top-level candidate traces (no ancestor in candidates) + is_sub_call = False + curr_parent = t.parent_trace_id + visited = set() + while curr_parent: + if curr_parent in visited: + break + visited.add(curr_parent) + if curr_parent in candidate_ids: + is_sub_call = True + break + parent_trace = next( + (x for x in traces if x.trace_id == curr_parent), + None, + ) + if parent_trace: + curr_parent = parent_trace.parent_trace_id + else: + break + if is_sub_call: + continue + + payload = t.payload or {} + tc_args = payload.get("args") or {} + + if isinstance(tc_args, dict): + filtered_args = { + k: v + for k, v in tc_args.items() + if k not in ("state", "tool_call_id") + } + else: + filtered_args = tc_args + + tc_result = payload.get("result") or payload.get("error") or "No result" + + interleaved_events.append( + { + "type": "tool_call", + "name": t.name, + "args": filtered_args, + "result": tc_result, + } + ) + + step_dict["interleaved_events"] = interleaved_events + step_dict["tool_calls"] = [ + e for e in interleaved_events if e["type"] == "tool_call" + ] + except Exception as e: + logger.error(f"Failed to retrieve traces for step {step.step_id}: {e}") + step_dict["interleaved_events"] = [] + step_dict["tool_calls"] = [] + + # Normalize physical coordinates to normalized ones for agent consumption + try: + extra = step_dict.get("extra_metadata") or {} + width = extra.get("width") or ( + getattr(self.ctx.device, "device_width", 1080) + if self.ctx and self.ctx.device + else 1080 + ) + height = extra.get("height") or ( + getattr(self.ctx.device, "device_height", 2400) + if self.ctx and self.ctx.device + else 2400 + ) + + if step_dict.get("action_taken"): + step_dict["action_taken"] = normalize_any_structure( + step_dict["action_taken"], width, height + ) + if step_dict.get("last_execution_result"): + step_dict["last_execution_result"] = normalize_any_structure( + step_dict["last_execution_result"], width, height + ) + except Exception as norm_err: + logger.error( + f"Failed to normalize coordinates in friendly steps history: {norm_err}" + ) + + result.append(step_dict) + return result + + def register_background_task(self, task_id: str, summary: str, trace_id: UUID | None = None): + if not self.current_session_id: + return + + record = BackgroundTaskRecord( + task_id=task_id, + session_id=self.current_session_id, + summary=summary, + status="running", + start_time=time.time(), + trace_id=str(trace_id) if trace_id else None, + ) + + if trace_id: + if not hasattr(self, "_bg_task_to_trace_id"): + self._bg_task_to_trace_id = {} + self._bg_task_to_trace_id[task_id] = str(trace_id) + + try: + task = asyncio.create_task( + asyncio.to_thread(self.storage.create_background_task, record) + ) + self._pending_tasks.add(task) + task.add_done_callback(self._pending_tasks.discard) + except RuntimeError: + + def run_and_cleanup(): + self.storage.create_background_task(record) + curr_thread = threading.current_thread() + if curr_thread in self._pending_threads: + self._pending_threads.remove(curr_thread) + + thread = threading.Thread(target=run_and_cleanup) + self._pending_threads.append(thread) + thread.start() + + self._publish("background_tasks_updated", self.get_all_background_tasks()) + + def unregister_background_task(self, task_id: str, status: str = "completed"): + end_time = time.time() + + trace_id_str = None + if hasattr(self, "_bg_task_to_trace_id"): + trace_id_str = self._bg_task_to_trace_id.pop(task_id, None) + + logs = "" + if trace_id_str and hasattr(self, "_accumulated_logs"): + logs = self._accumulated_logs.pop(trace_id_str, "") + + try: + task = asyncio.create_task( + asyncio.to_thread( + self.storage.update_background_task_status_and_logs, + task_id, + status, + end_time, + logs, + ) + ) + self._pending_tasks.add(task) + task.add_done_callback(self._pending_tasks.discard) + except RuntimeError: + + def run_and_cleanup(): + self.storage.update_background_task_status_and_logs(task_id, status, end_time, logs) + curr_thread = threading.current_thread() + if curr_thread in self._pending_threads: + self._pending_threads.remove(curr_thread) + + thread = threading.Thread(target=run_and_cleanup) + self._pending_threads.append(thread) + thread.start() + + self._publish("background_tasks_updated", self.get_all_background_tasks()) + + def get_all_background_tasks(self) -> list[dict]: + if not self.current_session_id: + return [] + records = self.storage.get_background_tasks(self.current_session_id) + return [r.model_dump() for r in records] + + def get_active_background_tasks(self) -> list[dict]: + return self.get_all_background_tasks() diff --git a/artemis/data_engine/models.py b/artemis/data_engine/models.py new file mode 100644 index 0000000..8e3989a --- /dev/null +++ b/artemis/data_engine/models.py @@ -0,0 +1,102 @@ +# 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 time +from typing import Any +from uuid import UUID +from pydantic import BaseModel, ConfigDict, Field + + +class SessionMetadata(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + session_id: UUID | str + initial_goal: str + start_time: float = Field(default_factory=time.time) + end_time: float | None = None + status: str = "running" # running, success, failed + device_info: dict[str, Any] = Field(default_factory=dict) + pid: int | None = None + video_filepath: str | None = None + + +class ImageRecord(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + image_name: str # SHA-256 hash + timestamp: float = Field(default_factory=time.time) + ocr_result: Any | None = None + ui_tree: Any | None = None + extra_metadata: dict[str, Any] = Field(default_factory=dict) + + +class StepRecord(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + step_id: UUID | str + session_id: UUID | str + step_number: int + timestamp: float = Field(default_factory=time.time) + pre_image_name: str | None = None + post_image_name: str | None = None + summary: str | None = None + action_taken: Any | None = None + operator_raw_thinking: str | None = None + operator_native_thinking: str | None = None + last_execution_result: Any | None = None + extra_metadata: dict[str, Any] = Field(default_factory=dict) + + +class TraceRecord(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + trace_id: UUID | str + session_id: UUID | str + step_id: UUID | str | None = None + parent_trace_id: UUID | str | None = None + type: str # "agent", "tool", "log" + name: str + timestamp: float = Field(default_factory=time.time) + duration: float | None = None + status: str = "success" # success, failed, running + payload: dict[str, Any] = Field(default_factory=dict) + + +class FailedOutputRecord(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + session_id: UUID | str + trace_id: UUID | str + model_name: str + prompt: str + raw_output: str + error_message: str + timestamp: float = Field(default_factory=time.time) + + +class BackgroundTaskRecord(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + task_id: str + session_id: UUID + summary: str + status: str = "running" # running, completed, failed + start_time: float = Field(default_factory=time.time) + end_time: float | None = None + trace_id: str | None = None + logs: str | None = None + + +class VideoRecordingRecord(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + video_id: UUID + session_id: UUID | None = None + device_id: str + start_time: float + end_time: float | None = None + local_video_path: str | None = None diff --git a/artemis/data_engine/storage.py b/artemis/data_engine/storage.py new file mode 100644 index 0000000..119135b --- /dev/null +++ b/artemis/data_engine/storage.py @@ -0,0 +1,1007 @@ +# 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. + +from collections.abc import Iterator +from contextlib import contextmanager +import difflib +import json +from pathlib import Path +import shutil +import sqlite3 +import threading +from typing import Any +import uuid +from uuid import UUID + +from artemis.data_engine.models import ( + BackgroundTaskRecord, + FailedOutputRecord, + ImageRecord, + SessionMetadata, + StepRecord, + TraceRecord, + VideoRecordingRecord, +) +from artemis.utils.logger import get_logger + + +def _safe_uuid(val: Any) -> UUID | str: + if val is None: + return "" + if isinstance(val, UUID): + return val + try: + return UUID(str(val)) + except (ValueError, TypeError): + return str(val) + + +class SafeJSONEncoder(json.JSONEncoder): + def default(self, obj): + if isinstance(obj, bytes): + return f"" + try: + return super().default(obj) + except TypeError: + return str(obj) + + +logger = get_logger(__name__) + + +class StorageManager: + """Manages persistence for Data Engine using SQLite and File System.""" + + def __init__(self, db_path: str | Path, base_trace_dir: str | Path): + self.db_path = Path(db_path) + self.base_trace_dir = Path(base_trace_dir) + self._lock = threading.RLock() + + # Ensure directories exist + self.db_path.parent.mkdir(parents=True, exist_ok=True) + self.base_trace_dir.mkdir(parents=True, exist_ok=True) + + self._init_db() + + @contextmanager + def _get_connection(self) -> Iterator[sqlite3.Connection]: + with self._lock: + conn = sqlite3.connect(self.db_path, timeout=30.0) + conn.row_factory = sqlite3.Row + try: + yield conn + finally: + conn.close() + + def _init_db(self): + """Initialize database tables.""" + with self._get_connection() as conn: + conn.execute("PRAGMA journal_mode=WAL") + conn.execute(""" + CREATE TABLE IF NOT EXISTS sessions ( + session_id TEXT PRIMARY KEY, + initial_goal TEXT, + start_time REAL, + end_time REAL, + status TEXT, + device_info TEXT, + video_filepath TEXT + ) + """) + + try: + conn.execute("ALTER TABLE sessions ADD COLUMN video_filepath TEXT") + except sqlite3.OperationalError: + pass + + try: + conn.execute("ALTER TABLE sessions ADD COLUMN pid INTEGER") + except sqlite3.OperationalError: + pass + + conn.execute(""" + CREATE TABLE IF NOT EXISTS images ( + image_name TEXT PRIMARY KEY, + timestamp REAL, + ocr_result TEXT, + ui_tree TEXT, + extra_metadata TEXT + ) + """) + + conn.execute(""" + CREATE TABLE IF NOT EXISTS steps ( + step_id TEXT PRIMARY KEY, + session_id TEXT, + step_number INTEGER, + timestamp REAL, + pre_image_name TEXT, + post_image_name TEXT, + summary TEXT, + action_taken TEXT, + operator_raw_thinking TEXT, + operator_native_thinking TEXT, + last_execution_result TEXT, + extra_metadata TEXT, + FOREIGN KEY(session_id) REFERENCES sessions(session_id), + FOREIGN KEY(pre_image_name) REFERENCES images(image_name), + FOREIGN KEY(post_image_name) REFERENCES images(image_name) + ) + """) + try: + conn.execute("ALTER TABLE steps ADD COLUMN action_taken TEXT") + except sqlite3.OperationalError: + pass + + try: + conn.execute("ALTER TABLE steps ADD COLUMN operator_raw_thinking TEXT") + except sqlite3.OperationalError: + pass + + try: + conn.execute("ALTER TABLE steps ADD COLUMN operator_native_thinking TEXT") + except sqlite3.OperationalError: + pass + + try: + conn.execute("ALTER TABLE steps ADD COLUMN last_execution_result TEXT") + except sqlite3.OperationalError: + pass + + conn.execute(""" + CREATE TABLE IF NOT EXISTS traces ( + trace_id TEXT PRIMARY KEY, + session_id TEXT, + step_id TEXT, + parent_trace_id TEXT, + type TEXT, + name TEXT, + timestamp REAL, + duration REAL, + status TEXT, + payload TEXT, + FOREIGN KEY(session_id) REFERENCES sessions(session_id), + FOREIGN KEY(step_id) REFERENCES steps(step_id), + FOREIGN KEY(parent_trace_id) REFERENCES traces(trace_id) + ) + """) + + conn.execute(""" + CREATE TABLE IF NOT EXISTS failed_outputs ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT, + trace_id TEXT, + model_name TEXT, + prompt TEXT, + raw_output TEXT, + error_message TEXT, + timestamp REAL, + FOREIGN KEY(session_id) REFERENCES sessions(session_id), + FOREIGN KEY(trace_id) REFERENCES traces(trace_id) + ) + """) + + conn.execute(""" + CREATE TABLE IF NOT EXISTS background_tasks ( + task_id TEXT PRIMARY KEY, + session_id TEXT, + summary TEXT, + status TEXT, + start_time REAL, + end_time REAL, + trace_id TEXT, + logs TEXT, + FOREIGN KEY(session_id) REFERENCES sessions(session_id) + ) + """) + conn.execute(""" + CREATE TABLE IF NOT EXISTS video_recordings ( + video_id TEXT PRIMARY KEY, + session_id TEXT, + device_id TEXT, + start_time REAL, + end_time REAL, + local_video_path TEXT, + FOREIGN KEY(session_id) REFERENCES sessions(session_id) + ) + """) + try: + conn.execute("ALTER TABLE background_tasks ADD COLUMN logs TEXT") + except sqlite3.OperationalError: + pass + conn.commit() + logger.info(f"Database initialized at {self.db_path}") + + def create_session(self, session: SessionMetadata): + """Insert a new session record.""" + with self._get_connection() as conn: + conn.execute( + """ + INSERT OR REPLACE INTO sessions (session_id, initial_goal, start_time, end_time, status, device_info, pid, video_filepath) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + str(session.session_id), + session.initial_goal, + session.start_time, + session.end_time, + session.status, + json.dumps(session.device_info), + session.pid, + session.video_filepath, + ), + ) + conn.commit() + + # Create session directory + session_dir = self.base_trace_dir / str(session.session_id) + session_dir.mkdir(parents=True, exist_ok=True) + logger.info(f"Session directory created at {session_dir}") + + def update_session(self, session: SessionMetadata): + """Update an existing session record.""" + with self._get_connection() as conn: + conn.execute( + """ + UPDATE sessions + SET end_time = ?, status = ?, device_info = ?, video_filepath = ? + WHERE session_id = ? + """, + ( + session.end_time, + session.status, + json.dumps(session.device_info), + session.video_filepath, + str(session.session_id), + ), + ) + conn.commit() + + def create_video_recording(self, record: VideoRecordingRecord): + """Insert a new video recording record.""" + with self._get_connection() as conn: + conn.execute( + """ + INSERT INTO video_recordings (video_id, session_id, device_id, start_time, end_time, local_video_path) + VALUES (?, ?, ?, ?, ?, ?) + """, + ( + str(record.video_id), + str(record.session_id) if record.session_id else None, + record.device_id, + record.start_time, + record.end_time, + record.local_video_path, + ), + ) + conn.commit() + + def update_video_recording(self, record: VideoRecordingRecord): + """Update an existing video recording record and session video path.""" + with self._get_connection() as conn: + conn.execute( + """ + UPDATE video_recordings + SET end_time = ?, local_video_path = ? + WHERE video_id = ? + """, + ( + record.end_time, + record.local_video_path, + str(record.video_id), + ), + ) + if record.session_id and record.local_video_path: + conn.execute( + """ + UPDATE sessions + SET video_filepath = ? + WHERE session_id = ? + """, + ( + record.local_video_path, + str(record.session_id), + ), + ) + conn.commit() + + def update_session_video_path(self, session_id: UUID, video_path: str): + """Update video filepath across session and video_recordings tables.""" + with self._get_connection() as conn: + conn.execute( + "UPDATE sessions SET video_filepath = ? WHERE session_id = ?", + (str(video_path), str(session_id)), + ) + conn.execute( + "UPDATE video_recordings SET local_video_path = ? WHERE session_id = ?", + (str(video_path), str(session_id)), + ) + conn.commit() + + def get_video_recording(self, video_id: UUID) -> VideoRecordingRecord | None: + """Retrieve a video recording record by ID.""" + with self._get_connection() as conn: + cursor = conn.execute( + "SELECT * FROM video_recordings WHERE video_id = ?", + (str(video_id),), + ) + row = cursor.fetchone() + if row: + return VideoRecordingRecord( + video_id=_safe_uuid(row["video_id"]), + session_id=_safe_uuid(row["session_id"]) if row["session_id"] else None, + device_id=row["device_id"], + start_time=row["start_time"], + end_time=row["end_time"], + local_video_path=row["local_video_path"], + ) + return None + + def create_image(self, image: ImageRecord): + """Insert a new image record.""" + with self._get_connection() as conn: + conn.execute( + """ + INSERT OR IGNORE INTO images (image_name, timestamp, ocr_result, ui_tree, extra_metadata) + VALUES (?, ?, ?, ?, ?) + """, + ( + image.image_name, + image.timestamp, + json.dumps(image.ocr_result) if image.ocr_result else None, + json.dumps(image.ui_tree) if image.ui_tree else None, + json.dumps(image.extra_metadata), + ), + ) + conn.commit() + + def get_image(self, image_name: str) -> ImageRecord | None: + """Retrieve an image by its name.""" + with self._get_connection() as conn: + cursor = conn.execute("SELECT * FROM images WHERE image_name = ?", (image_name,)) + row = cursor.fetchone() + if row: + return ImageRecord( + image_name=row["image_name"], + timestamp=row["timestamp"], + ocr_result=json.loads(row["ocr_result"]) if row["ocr_result"] else None, + ui_tree=json.loads(row["ui_tree"]) if row["ui_tree"] else None, + extra_metadata=json.loads(row["extra_metadata"]), + ) + return None + + def update_image_data(self, image_name: str, ocr_result: Any | None, ui_tree: Any | None): + """Update image record with OCR and UI tree if they are not already set.""" + with self._get_connection() as conn: + ocr_json = json.dumps(ocr_result) if ocr_result is not None else None + ui_tree_json = json.dumps(ui_tree) if ui_tree is not None else None + + conn.execute( + """ + UPDATE images + SET ocr_result = COALESCE(ocr_result, ?), + ui_tree = COALESCE(ui_tree, ?) + WHERE image_name = ? + """, + (ocr_json, ui_tree_json, image_name), + ) + conn.commit() + + def search_ui_by_hash(self, image_name: str, query: str, threshold: float = 0.6) -> list[dict]: + """Search for text in UI data (XML and OCR) retrieved by image hash.""" + record = self.get_image(image_name) + if not record: + return [] + + matches = [] + + def check_match(text, target_query): + if not text: + return 0.0 + return difflib.SequenceMatcher(None, target_query.lower(), text.lower()).ratio() + + # Search in ui_tree (raw XML elements) + if record.ui_tree: + for node in record.ui_tree: + best_ratio = 0.0 + best_target = "" + + # Check text + text = node.get("text") + if text: + ratio = check_match(text, query) + if ratio > best_ratio: + best_ratio = ratio + best_target = text + + # Check content-desc + desc = node.get("content-desc") + if desc: + ratio = check_match(desc, query) + if ratio > best_ratio: + best_ratio = ratio + best_target = desc + + if best_ratio > threshold: + matches.append( + { + "score": best_ratio, + "type": "xml", + "node": node, + "matched_text": best_target, + } + ) + + # Search in ocr_result (raw OCR results) + if record.ocr_result: + for ocr in record.ocr_result: + text = ocr.get("text") + if text: + ratio = check_match(text, query) + if ratio > threshold: + matches.append( + { + "score": ratio, + "type": "ocr", + "node": ocr, + "matched_text": text, + } + ) + + # Sort by score descending + matches.sort(key=lambda x: x["score"], reverse=True) + return matches + + def create_step(self, step: StepRecord): + """Insert a new step record.""" + with self._get_connection() as conn: + conn.execute( + """ + INSERT INTO steps ( + step_id, session_id, step_number, timestamp, + pre_image_name, post_image_name, + summary, action_taken, operator_raw_thinking, operator_native_thinking, last_execution_result, extra_metadata + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + str(step.step_id), + str(step.session_id), + step.step_number, + step.timestamp, + step.pre_image_name, + step.post_image_name, + step.summary, + json.dumps(step.action_taken) if step.action_taken else None, + step.operator_raw_thinking, + step.operator_native_thinking, + json.dumps(step.last_execution_result) if step.last_execution_result else None, + json.dumps(step.extra_metadata), + ), + ) + conn.commit() + + def update_step_action(self, step_id: UUID, action_taken: dict): + """Update the action taken for a step.""" + with self._get_connection() as conn: + conn.execute( + """ + UPDATE steps + SET action_taken = ? + WHERE step_id = ? + """, + ( + json.dumps(action_taken), + str(step_id), + ), + ) + conn.commit() + + def update_step_summary(self, step_id: UUID, summary: str): + """Update the summary for a step.""" + with self._get_connection() as conn: + conn.execute( + """ + UPDATE steps + SET summary = ? + WHERE step_id = ? + """, + ( + summary, + str(step_id), + ), + ) + conn.commit() + + def update_step_thinking(self, step_id: UUID, operator_raw_thinking: str): + """Update the raw thinking process for a step.""" + with self._get_connection() as conn: + conn.execute( + """ + UPDATE steps + SET operator_raw_thinking = ? + WHERE step_id = ? + """, + ( + operator_raw_thinking, + str(step_id), + ), + ) + conn.commit() + + def update_step_native_thinking(self, step_id: UUID, operator_native_thinking: str): + """Update the native thinking process for a step.""" + with self._get_connection() as conn: + conn.execute( + """ + UPDATE steps + SET operator_native_thinking = ? + WHERE step_id = ? + """, + ( + operator_native_thinking, + str(step_id), + ), + ) + conn.commit() + + def update_step_execution_result( + self, + step_id: UUID, + last_execution_result: dict, + post_image_name: str | None = None, + ): + """Update the validator execution result and post_image_name for a step.""" + with self._get_connection() as conn: + if post_image_name: + cursor = conn.execute( + "SELECT pre_image_name FROM steps WHERE step_id = ?", + (str(step_id),), + ) + row = cursor.fetchone() + if row and row["pre_image_name"] == post_image_name: + post_image_name = None + + if post_image_name: + conn.execute( + """ + UPDATE steps + SET last_execution_result = ?, post_image_name = ? + WHERE step_id = ? + """, + ( + json.dumps(last_execution_result), + post_image_name, + str(step_id), + ), + ) + else: + conn.execute( + """ + UPDATE steps + SET last_execution_result = ? + WHERE step_id = ? + """, + ( + json.dumps(last_execution_result), + str(step_id), + ), + ) + conn.commit() + + def get_session(self, session_id: UUID | str) -> SessionMetadata | None: + """Retrieve a session by ID.""" + with self._get_connection() as conn: + cursor = conn.execute( + "SELECT * FROM sessions WHERE session_id = ?", + (str(session_id),), + ) + row = cursor.fetchone() + if row: + row_dict = dict(row) + return SessionMetadata( + session_id=_safe_uuid(row_dict["session_id"]), + initial_goal=row_dict["initial_goal"], + start_time=row_dict["start_time"], + end_time=row_dict["end_time"], + status=row_dict["status"], + device_info=json.loads(row_dict["device_info"]), + pid=row_dict.get("pid"), + video_filepath=row_dict.get("video_filepath"), + ) + return None + + def get_steps(self, session_id: UUID | str) -> list[StepRecord]: + """Retrieve all steps for a session.""" + steps = [] + with self._get_connection() as conn: + cursor = conn.execute( + "SELECT * FROM steps WHERE session_id = ? ORDER BY step_number ASC", + (str(session_id),), + ) + for row in cursor.fetchall(): + action_taken = None + if "action_taken" in row.keys() and row["action_taken"]: + try: + action_taken = json.loads(row["action_taken"]) + except Exception: + pass + + operator_raw_thinking = None + if "operator_raw_thinking" in row.keys() and row["operator_raw_thinking"]: + operator_raw_thinking = row["operator_raw_thinking"] + + operator_native_thinking = None + if "operator_native_thinking" in row.keys() and row["operator_native_thinking"]: + operator_native_thinking = row["operator_native_thinking"] + + last_execution_result = None + if "last_execution_result" in row.keys() and row["last_execution_result"]: + try: + last_execution_result = json.loads(row["last_execution_result"]) + except Exception: + pass + + steps.append( + StepRecord( + step_id=_safe_uuid(row["step_id"]), + session_id=_safe_uuid(row["session_id"]), + step_number=row["step_number"], + timestamp=row["timestamp"], + pre_image_name=row["pre_image_name"], + post_image_name=row["post_image_name"], + summary=row["summary"], + action_taken=action_taken, + operator_raw_thinking=operator_raw_thinking, + operator_native_thinking=operator_native_thinking, + last_execution_result=last_execution_result, + extra_metadata=json.loads(row["extra_metadata"]), + ) + ) + return steps + + def get_traces_for_step(self, step_id: UUID | str) -> list[TraceRecord]: + """Retrieve all traces for a step.""" + traces = [] + with self._get_connection() as conn: + cursor = conn.execute( + "SELECT * FROM traces WHERE step_id = ? ORDER BY timestamp ASC", + (str(step_id),), + ) + for row in cursor.fetchall(): + traces.append( + TraceRecord( + trace_id=_safe_uuid(row["trace_id"]), + session_id=_safe_uuid(row["session_id"]), + step_id=_safe_uuid(row["step_id"]) if row["step_id"] else None, + parent_trace_id=_safe_uuid(row["parent_trace_id"]) + if row["parent_trace_id"] + else None, + type=row["type"], + name=row["name"], + timestamp=row["timestamp"], + duration=row["duration"], + status=row["status"], + payload=json.loads(row["payload"]) if row["payload"] else {}, + ) + ) + return traces + + def get_steps_with_traces(self, session_id: UUID) -> list[tuple[StepRecord, list[TraceRecord]]]: + """Retrieve all steps and their associated traces for a session.""" + steps = self.get_steps(session_id) + result = [] + for step in steps: + traces = self.get_traces_for_step(step.step_id) + result.append((step, traces)) + return result + + def create_trace(self, trace: TraceRecord): + """Insert or update a trace record (handles out-of-order writes and preserves start timestamps).""" + with self._get_connection() as conn: + conn.execute( + """ + INSERT INTO traces ( + trace_id, session_id, step_id, parent_trace_id, type, name, + timestamp, duration, status, payload + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(trace_id) DO UPDATE SET + status = CASE + WHEN excluded.status = 'running' AND traces.status IN ('success', 'failed') THEN traces.status + ELSE excluded.status + END, + duration = COALESCE(excluded.duration, traces.duration), + timestamp = MIN(traces.timestamp, excluded.timestamp), + payload = CASE + WHEN excluded.status = 'running' AND traces.status IN ('success', 'failed') THEN traces.payload + ELSE excluded.payload + END + """, + ( + str(trace.trace_id), + str(trace.session_id), + str(trace.step_id) if trace.step_id else None, + str(trace.parent_trace_id) if trace.parent_trace_id else None, + trace.type, + trace.name, + trace.timestamp, + trace.duration, + trace.status, + json.dumps(trace.payload, cls=SafeJSONEncoder), + ), + ) + conn.commit() + + def create_failed_output(self, record: FailedOutputRecord): + """Insert a new failed output record (crime scene).""" + with self._get_connection() as conn: + conn.execute( + """ + INSERT INTO failed_outputs ( + session_id, trace_id, model_name, prompt, raw_output, error_message, timestamp + ) + VALUES (?, ?, ?, ?, ?, ?, ?) + """, + ( + str(record.session_id), + str(record.trace_id), + record.model_name, + record.prompt, + record.raw_output, + record.error_message, + record.timestamp, + ), + ) + conn.commit() + + def create_background_task(self, record: BackgroundTaskRecord): + with self._get_connection() as conn: + conn.execute( + """ + INSERT OR REPLACE INTO background_tasks ( + task_id, session_id, summary, status, start_time, end_time, trace_id, logs + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + record.task_id, + str(record.session_id), + record.summary, + record.status, + record.start_time, + record.end_time, + record.trace_id, + record.logs, + ), + ) + conn.commit() + + def update_background_task_status(self, task_id: str, status: str, end_time: float): + with self._get_connection() as conn: + conn.execute( + """ + UPDATE background_tasks + SET status = ?, end_time = ? + WHERE task_id = ? + """, + (status, end_time, task_id), + ) + conn.commit() + + def update_background_task_status_and_logs( + self, task_id: str, status: str, end_time: float, logs: str + ): + with self._get_connection() as conn: + conn.execute( + """ + UPDATE background_tasks + SET status = ?, end_time = ?, logs = ? + WHERE task_id = ? + """, + (status, end_time, logs, task_id), + ) + conn.commit() + + def get_background_tasks(self, session_id: UUID) -> list[BackgroundTaskRecord]: + tasks = [] + with self._get_connection() as conn: + cursor = conn.execute( + "SELECT * FROM background_tasks WHERE session_id = ? ORDER BY start_time ASC", + (str(session_id),), + ) + for row in cursor.fetchall(): + logs = row["logs"] if "logs" in row.keys() else None + tasks.append( + BackgroundTaskRecord( + task_id=row["task_id"], + session_id=_safe_uuid(row["session_id"]), + summary=row["summary"], + status=row["status"], + start_time=row["start_time"], + end_time=row["end_time"], + trace_id=row["trace_id"], + logs=logs, + ) + ) + return tasks + + def get_step_traces(self, step_id: UUID | str) -> list[TraceRecord]: + """Retrieve all traces for a specific step.""" + traces = [] + with self._get_connection() as conn: + cursor = conn.execute( + "SELECT * FROM traces WHERE step_id = ? ORDER BY timestamp ASC", + (str(step_id),), + ) + for row in cursor.fetchall(): + traces.append( + TraceRecord( + trace_id=_safe_uuid(row["trace_id"]), + session_id=_safe_uuid(row["session_id"]), + step_id=_safe_uuid(row["step_id"]) if row["step_id"] else None, + parent_trace_id=_safe_uuid(row["parent_trace_id"]) + if row["parent_trace_id"] + else None, + type=row["type"], + name=row["name"], + timestamp=row["timestamp"], + duration=row["duration"], + status=row["status"], + payload=json.loads(row["payload"]) if row["payload"] else {}, + ) + ) + return traces + + def clear_all_data(self): + """Clear all data from SQLite tables and delete session directories.""" + with self._get_connection() as conn: + conn.execute("PRAGMA foreign_keys = OFF") + tables = [ + "failed_outputs", + "traces", + "background_tasks", + "steps", + "images", + "sessions", + ] + for table in tables: + try: + conn.execute(f"DELETE FROM {table}") + except sqlite3.OperationalError as e: + logger.warning(f"Failed to clear table {table}: {e}") + conn.commit() + conn.execute("VACUUM") + conn.commit() + logger.info("Database tables cleared and vacuumed.") + + # Delete session directories + if self.base_trace_dir.exists(): + for path in self.base_trace_dir.iterdir(): + if path.is_dir(): + if path.name == "images": + # Clear images + for img_path in path.iterdir(): + if img_path.is_file(): + try: + img_path.unlink() + except Exception as e: + logger.error(f"Failed to delete image {img_path}: {e}") + elif path.name.startswith("web_") or self._is_uuid(path.name): + try: + shutil.rmtree(path) + except Exception as e: + logger.error(f"Failed to delete directory {path}: {e}") + elif path.is_file() and path.name == "mcp_server.log": + try: + path.unlink() + except Exception as e: + logger.error(f"Failed to delete log file {path}: {e}") + + def _is_uuid(self, val: str) -> bool: + # Support both pure UUID and UUID with suffix (e.g., UUID_PASS_...) + parts = val.split("_") + target = parts[0] if parts else val + try: + uuid.UUID(target) + return True + except ValueError: + return False + + def delete_session(self, session_id: UUID): + """Delete all data associated with a session, including files on disk.""" + session_id_str = str(session_id) + + # 1. Get video paths before deleting from DB + video_paths = [] + try: + with self._get_connection() as conn: + cursor = conn.execute( + "SELECT local_video_path FROM video_recordings WHERE session_id = ?", + (session_id_str,), + ) + for row in cursor.fetchall(): + if row["local_video_path"]: + video_paths.append(Path(row["local_video_path"])) + except sqlite3.OperationalError as e: + logger.warning(f"Could not query video_recordings: {e}") + + # 2. Delete from DB tables + with self._get_connection() as conn: + conn.execute("PRAGMA foreign_keys = OFF") + tables = [ + "failed_outputs", + "traces", + "background_tasks", + "video_recordings", + "steps", + "sessions", + ] + for table in tables: + try: + conn.execute( + f"DELETE FROM {table} WHERE session_id = ?", + (session_id_str,), + ) + except sqlite3.OperationalError as e: + logger.warning(f"Failed to delete from table {table}: {e}") + conn.commit() + conn.execute("VACUUM") + conn.commit() + + logger.info(f"Database records for session {session_id} cleared.") + + # 3. Delete session directory (notes, etc.) + session_dir = self.base_trace_dir / session_id_str + if session_dir.exists() and session_dir.is_dir(): + try: + shutil.rmtree(session_dir) + logger.info(f"Deleted session directory: {session_dir}") + except Exception as e: + logger.error(f"Failed to delete session directory {session_dir}: {e}") + + # 4. Delete video files and their containing folders (which might have been renamed) + for video_path in video_paths: + # video_path is likely: /path/to/traces/task_name/recording.mp4 + # If it was renamed, it might be: /path/to/traces/task_name_PASS_timestamp/recording.mp4 + + task_name = video_path.parent.name + + if task_name and task_name != "traces" and task_name != "..": + # Search for directories starting with task_name in base_trace_dir + try: + for path in self.base_trace_dir.iterdir(): + if path.is_dir() and path.name.startswith(task_name): + try: + shutil.rmtree(path) + logger.info(f"Deleted trace/video directory: {path}") + except Exception as e: + logger.error(f"Failed to delete directory {path}: {e}") + elif path.is_file() and path.name.startswith(task_name): + try: + path.unlink() + logger.info(f"Deleted trace/video file: {path}") + except Exception as e: + logger.error(f"Failed to delete file {path}: {e}") + except Exception as e: + logger.error(f"Error iterating base_trace_dir for cleanup: {e}") + + if video_path.exists(): + try: + video_path.unlink() + logger.info(f"Deleted video file directly: {video_path}") + if video_path.parent.exists() and not any(video_path.parent.iterdir()): + video_path.parent.rmdir() + except Exception as e: + logger.error(f"Failed to delete video file directly {video_path}: {e}") diff --git a/artemis/data_engine/trace.py b/artemis/data_engine/trace.py new file mode 100644 index 0000000..d0cda37 --- /dev/null +++ b/artemis/data_engine/trace.py @@ -0,0 +1,675 @@ +# 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 +from collections.abc import Callable +import functools +import hashlib +import inspect +import json +import time +from typing import Any +import uuid +from uuid import UUID + +from langchain_core.callbacks import BaseCallbackHandler +from langchain_core.messages import BaseMessage +from langchain_core.tools import BaseTool + +from artemis.context import ArtemisContext +from artemis.data_engine.context_vars import ( + CURRENT_NODE_NAME, + CURRENT_TRACE_ID, +) + + +def smart_serialize(obj: Any) -> str: + """Smartly serializes an object to string, handling complex types.""" + + if obj.__class__.__name__ == "State": + return "" + + if isinstance(obj, (dict, list)): + try: + cleaned = _deep_smart_clean(obj) + return json.dumps(cleaned, ensure_ascii=False) + except Exception: + return str(obj) + elif isinstance(obj, bytes): + return f"" + elif isinstance(obj, str): + if len(obj) > 5000: + clean_str = obj + if "base64," in clean_str: + clean_str = clean_str.split("base64,")[-1] + + try: + img_bytes = base64.b64decode(clean_str, validate=False) + if len(img_bytes) > 1000: + sha = hashlib.sha256(img_bytes).hexdigest() + return f"" + except Exception: + pass + return ( + f"" + ) + return obj + else: + return str(obj) + + +def _deep_smart_clean(obj: Any) -> Any: + """Recursively cleans dictionaries and lists using smart_serialize rules for strings and bytes.""" + if isinstance(obj, dict): + return {k: _deep_smart_clean(v) for k, v in obj.items()} + elif isinstance(obj, list): + return [_deep_smart_clean(v) for v in obj] + elif isinstance(obj, bytes): + return f"" + elif isinstance(obj, str): + if len(obj) > 5000: + clean_str = obj + if "base64," in clean_str: + clean_str = clean_str.split("base64,")[-1] + + try: + img_bytes = base64.b64decode(clean_str, validate=False) + if len(img_bytes) > 1000: + sha = hashlib.sha256(img_bytes).hexdigest() + return f"" + except Exception: + pass + return ( + f"" + ) + return obj + else: + return obj + + +def serialize_error(error: Any) -> str: + """Returns a user-friendly and informative string representation of an exception or error message.""" + if not error: + return "" + if isinstance(error, str): + return error + err_str = str(error) + if not err_str: + return f"{error.__class__.__name__}" + return f"{error.__class__.__name__}: {err_str}" + + +class DataEngineCallbackHandler(BaseCallbackHandler): + """Callback handler to record executions to DataEngine.""" + + def __init__(self, ctx: ArtemisContext): + self.ctx = ctx + self.running_spans = {} + self.chain_parents = {} + + def _resolve_parent_id( + self, parent_run_id: uuid.UUID | None, kwargs: dict[str, Any] = None + ) -> uuid.UUID | None: + curr = parent_run_id + while curr in self.chain_parents: + curr = self.chain_parents[curr] + if curr: + return curr + + # Check in metadata passed via kwargs + if kwargs and "metadata" in kwargs: + metadata = kwargs["metadata"] or {} + p_id = metadata.get("parent_trace_id") + if p_id: + try: + return uuid.UUID(p_id) + except ValueError: + pass + + return CURRENT_TRACE_ID.get() + + def on_chain_start( + self, + serialized: dict[str, Any], + inputs: dict[str, Any], + *, + run_id: uuid.UUID, + parent_run_id: uuid.UUID | None = None, + **kwargs: Any, + ) -> Any: + self.chain_parents[run_id] = parent_run_id or CURRENT_TRACE_ID.get() + + def on_chain_end(self, outputs: dict[str, Any], *, run_id: uuid.UUID, **kwargs: Any) -> Any: + self.chain_parents.pop(run_id, None) + + def on_chain_error(self, error: Any, *, run_id: uuid.UUID, **kwargs: Any) -> Any: + self.chain_parents.pop(run_id, None) + + def on_tool_start( + self, + serialized: dict[str, Any], + input_str: str, + *, + run_id: uuid.UUID, + parent_run_id: uuid.UUID | None = None, + **kwargs: Any, + ) -> Any: + pass + + def on_tool_end(self, output: str, *, run_id: uuid.UUID, **kwargs: Any) -> Any: + pass + + def on_tool_error(self, error: Any, *, run_id: uuid.UUID, **kwargs: Any) -> Any: + pass + + def on_llm_start( + self, + serialized: dict[str, Any], + prompts: list[str], + *, + run_id: uuid.UUID, + parent_run_id: uuid.UUID | None = None, + **kwargs: Any, + ) -> Any: + """Run when LLM starts running (completion API).""" + if self.ctx.data_engine: + name = serialized.get("name") or "llm" + self.running_spans[run_id] = { + "prompts": prompts, + "name": name, + "start_time": time.time(), + } + + self.ctx.data_engine.record_trace( + type="llm_call", + name=name, + payload={"prompts": prompts}, + trace_id=run_id, + parent_trace_id=self._resolve_parent_id(parent_run_id, kwargs), + status="running", + step_id=self.ctx.data_engine.current_step_id, + ) + + def on_chat_model_start( + self, + serialized: dict[str, Any], + messages: list[list[BaseMessage]], + *, + run_id: uuid.UUID, + parent_run_id: uuid.UUID | None = None, + **kwargs: Any, + ) -> Any: + """Run when Chat Model starts running.""" + if self.ctx.data_engine: + flat_messages = [] + if messages: + for msg in messages[0]: + flat_messages.append(self._serialize_message(msg)) + + name = serialized.get("name") or "llm" + self.running_spans[run_id] = { + "messages": flat_messages, + "name": name, + "start_time": time.time(), + } + + self.ctx.data_engine.record_trace( + type="llm_call", + name=name, + payload={"messages": flat_messages}, + trace_id=run_id, + parent_trace_id=self._resolve_parent_id(parent_run_id, kwargs), + status="running", + step_id=self.ctx.data_engine.current_step_id, + ) + + def _safe_uuid5(self, parent_id, name: str) -> uuid.UUID: + try: + pid = parent_id if isinstance(parent_id, uuid.UUID) else uuid.UUID(str(parent_id)) + return uuid.uuid5(pid, name) + except Exception: + return uuid.uuid4() + + def on_llm_end( + self, + response: Any, + *, + run_id: uuid.UUID, + parent_run_id: uuid.UUID | None = None, + **kwargs: Any, + ) -> Any: + """Run when LLM ends running.""" + if self.ctx.data_engine: + generations = getattr(response, "generations", []) + flat_generations = [] + if generations: + for gen in generations[0]: + message = getattr(gen, "message", None) + if message: + flat_generations.append(self._serialize_message(message)) + else: + flat_generations.append({"text": getattr(gen, "text", "")}) + + span_data = self.running_spans.pop(run_id, {}) + messages = span_data.get("messages", []) + prompts = span_data.get("prompts", []) + name = span_data.get("name", "llm") + start_time = span_data.get("start_time") + duration = time.time() - start_time if start_time else None + + payload = {"response": flat_generations} + if messages: + payload["messages"] = messages + if prompts: + payload["prompts"] = prompts + + # Extract token usage if available from LLM output + usage_data = None + if hasattr(response, "llm_output") and isinstance(response.llm_output, dict): + usage_data = response.llm_output.get("token_usage") or response.llm_output.get( + "usage_metadata" + ) + if not usage_data and generations: + for gen in generations[0]: + msg = getattr(gen, "message", None) + if msg: + if hasattr(msg, "usage_metadata") and msg.usage_metadata: + usage_data = msg.usage_metadata + break + if hasattr(msg, "response_metadata") and isinstance( + msg.response_metadata, dict + ): + usage_data = msg.response_metadata.get( + "usage_metadata" + ) or msg.response_metadata.get("token_usage") + if usage_data: + break + if usage_data: + payload["token_usage"] = usage_data + + self.ctx.data_engine.record_trace( + type="llm_call", + name=name, + payload=payload, + trace_id=run_id, + parent_trace_id=self._resolve_parent_id(parent_run_id), + status="success", + duration=duration, + step_id=self.ctx.data_engine.current_step_id, + ) + + native_thoughts = [] + raw_thoughts = [] + for gen_item in flat_generations: + content = gen_item.get("content") + if isinstance(content, list): + for block in content: + if isinstance(block, dict): + if block.get("type") == "thinking" and block.get("thinking"): + native_thoughts.append(str(block.get("thinking")).strip()) + elif block.get("type") == "text" and block.get("text"): + raw_thoughts.append(str(block.get("text")).strip()) + elif isinstance(content, str) and content.strip(): + raw_thoughts.append(content.strip()) + elif gen_item.get("text") and str(gen_item.get("text")).strip(): + raw_thoughts.append(str(gen_item.get("text")).strip()) + + if native_thoughts: + self.ctx.data_engine.record_trace( + type="thinking", + name="thinking", + payload={"thought": "\n\n".join(native_thoughts)}, + trace_id=self._safe_uuid5(run_id, "thinking"), + parent_trace_id=run_id, + status="success", + step_id=self.ctx.data_engine.current_step_id, + ) + if raw_thoughts: + self.ctx.data_engine.record_trace( + type="raw_thinking", + name="raw_thinking", + payload={"thought": "\n\n".join(raw_thoughts)}, + trace_id=self._safe_uuid5(run_id, "raw_thinking"), + parent_trace_id=run_id, + status="success", + step_id=self.ctx.data_engine.current_step_id, + ) + + def on_llm_error( + self, + error: Any, + *, + run_id: uuid.UUID, + parent_run_id: uuid.UUID | None = None, + **kwargs: Any, + ) -> Any: + """Run when LLM errors.""" + if self.ctx.data_engine: + span_data = self.running_spans.pop(run_id, {}) + messages = span_data.get("messages", []) + prompts = span_data.get("prompts", []) + name = span_data.get("name", "llm") + start_time = span_data.get("start_time") + duration = time.time() - start_time if start_time else None + + payload = {"error": serialize_error(error)} + if messages: + payload["messages"] = messages + if prompts: + payload["prompts"] = prompts + + self.ctx.data_engine.record_trace( + type="llm_call", + name=name, + payload=payload, + trace_id=run_id, + parent_trace_id=self._resolve_parent_id(parent_run_id), + status="failed", + duration=duration, + step_id=self.ctx.data_engine.current_step_id, + ) + + def _serialize_message(self, msg: BaseMessage) -> dict[str, Any]: + return { + "type": msg.type, + "content": msg.content, + "additional_kwargs": msg.additional_kwargs, + } + + +def trace( + type: str, + name: str | None = None, + ctx: ArtemisContext | None = None, + serializer: Callable[[Any], dict[str, Any]] | None = None, +): + """Decorator to trace execution of tools and agents in Data Engine. + + Args: + type: "agent" or "tool" + name: Optional name of the trace. Defaults to function name. + ctx: Optional ArtemisContext. If not provided, tries to find it in + arguments. + serializer: Optional function to customize how result is stored. + """ + + def decorator(func): + trace_name = name or func.__name__ + sig = inspect.signature(func) + + def _get_ctx(*args, **kwargs) -> ArtemisContext | None: + if ctx: + return ctx + # Try to find in args + for arg in args: + if isinstance(arg, ArtemisContext): + return arg + if hasattr(arg, "ctx") and isinstance(getattr(arg, "ctx"), ArtemisContext): + return getattr(arg, "ctx") + # Try to find in kwargs + for v in kwargs.values(): + if isinstance(v, ArtemisContext): + return v + if hasattr(v, "ctx") and isinstance(getattr(v, "ctx"), ArtemisContext): + return getattr(v, "ctx") + return None + + def _record( + ctx_obj, + args, + kwargs, + result, + status, + duration, + trace_id, + parent_trace_id, + error=None, + ): + if not ctx_obj or not ctx_obj.data_engine: + return + + try: + bound = sig.bind(*args, **kwargs) + filtered_args = {} + for k, v in bound.arguments.items(): + if k not in ("self", "ctx", "state", "tool_call_id"): + filtered_args[k] = smart_serialize(v) + payload = {"args": filtered_args} + except Exception: + filtered_kwargs = { + k: v for k, v in kwargs.items() if k not in ("state", "tool_call_id") + } + payload = {"args": {k: smart_serialize(v) for k, v in filtered_kwargs.items()}} + if args: + payload["positional_args"] = [smart_serialize(a) for a in args] + + if status == "success": + if serializer: + payload["result"] = serializer(result) + else: + payload["result"] = smart_serialize(result) + elif status == "failed": + payload["error"] = serialize_error(error) + + step_id = None + if ( + hasattr(ctx_obj.data_engine, "current_step_id") + and ctx_obj.data_engine.current_step_id + ): + step_id = ctx_obj.data_engine.current_step_id + else: + try: + bound = sig.bind(*args, **kwargs) + state_arg = bound.arguments.get("state") + if ( + state_arg + and hasattr(state_arg, "current_step_id") + and state_arg.current_step_id + ): + step_id = UUID(str(state_arg.current_step_id)) + except Exception: + pass + + ctx_obj.data_engine.record_trace( + type=type, + name=trace_name, + payload=payload, + status=status, + duration=duration, + trace_id=trace_id, + parent_trace_id=parent_trace_id, + step_id=step_id, + ) + + @functools.wraps(func) + async def async_wrapper(*args, **kwargs): + ctx_obj = _get_ctx(*args, **kwargs) + start_time = time.time() + + trace_id = uuid.uuid4() + parent_id = CURRENT_TRACE_ID.get() + + token = CURRENT_TRACE_ID.set(trace_id) + token_name = CURRENT_NODE_NAME.set(trace_name) + + # Record trace at start with 'running' status + _record(ctx_obj, args, kwargs, None, "running", 0.0, trace_id, parent_id) + + try: + result = await func(*args, **kwargs) + duration = time.time() - start_time + _record( + ctx_obj, + args, + kwargs, + result, + "success", + duration, + trace_id, + parent_id, + ) + return result + except Exception as e: + duration = time.time() - start_time + _record( + ctx_obj, + args, + kwargs, + None, + "failed", + duration, + trace_id, + parent_id, + error=e, + ) + raise e + finally: + CURRENT_TRACE_ID.reset(token) + CURRENT_NODE_NAME.reset(token_name) + + @functools.wraps(func) + def sync_wrapper(*args, **kwargs): + ctx_obj = _get_ctx(*args, **kwargs) + start_time = time.time() + + trace_id = uuid.uuid4() + parent_id = CURRENT_TRACE_ID.get() + + token = CURRENT_TRACE_ID.set(trace_id) + token_name = CURRENT_NODE_NAME.set(trace_name) + + # Record trace at start with 'running' status + _record(ctx_obj, args, kwargs, None, "running", 0.0, trace_id, parent_id) + + try: + result = func(*args, **kwargs) + duration = time.time() - start_time + _record( + ctx_obj, + args, + kwargs, + result, + "success", + duration, + trace_id, + parent_id, + ) + return result + except Exception as e: + duration = time.time() - start_time + _record( + ctx_obj, + args, + kwargs, + None, + "failed", + duration, + trace_id, + parent_id, + error=e, + ) + raise e + finally: + CURRENT_TRACE_ID.reset(token) + CURRENT_NODE_NAME.reset(token_name) + + if inspect.iscoroutinefunction(func): + return async_wrapper + else: + return sync_wrapper + + return decorator + + +class TraceSpan: + """Context manager to create a child trace span with dynamic payload and proactive state.""" + + def __init__( + self, + name: str, + trace_type: str = "span", + ctx: ArtemisContext | None = None, + serializer: Callable[[Any], dict[str, Any]] | None = None, + ): + self.name = name + self.trace_type = trace_type + self.ctx = ctx + self.serializer = serializer + self.start_time = None + self.trace_id = uuid.uuid4() + self.parent_id = None + self.token = None + self.status = "success" + self.result = None + self.error = None + self.payload = {} + + def __enter__(self): + self.start_time = time.time() + self.parent_id = CURRENT_TRACE_ID.get() + self.token = CURRENT_TRACE_ID.set(self.trace_id) + + if self.ctx and self.ctx.data_engine: + step_id = getattr(self.ctx.data_engine, "current_step_id", None) + self.ctx.data_engine.record_trace( + type=self.trace_type, + name=self.name, + payload={"status": "running", **self.payload}, + status="running", + parent_trace_id=self.parent_id, + step_id=step_id, + trace_id=self.trace_id, + ) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + duration = time.time() - self.start_time + CURRENT_TRACE_ID.reset(self.token) + + if exc_type: + self.status = "failed" + self.error = exc_val + + if self.ctx and self.ctx.data_engine: + final_payload = dict(self.payload) + if self.status == "success": + if self.serializer and self.result: + final_payload["result"] = self.serializer(self.result) + elif self.result: + final_payload["result"] = smart_serialize(self.result) + else: + final_payload["error"] = serialize_error(self.error) + + step_id = getattr(self.ctx.data_engine, "current_step_id", None) + self.ctx.data_engine.record_trace( + type=self.trace_type, + name=self.name, + payload=final_payload, + status=self.status, + duration=duration, + trace_id=self.trace_id, + parent_trace_id=self.parent_id, + step_id=step_id, + ) + + +def trace_langchain_tool(tool: BaseTool, ctx: ArtemisContext) -> BaseTool: + """Wraps a LangChain tool to trace its execution. + + [REFACTORED] Tracing is now completely delegated to + DataEngineCallbackHandler infrastructure. + """ + return tool diff --git a/artemis/drivers/__init__.py b/artemis/drivers/__init__.py new file mode 100644 index 0000000..c240b94 --- /dev/null +++ b/artemis/drivers/__init__.py @@ -0,0 +1,48 @@ +# 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. + +"""Universal Device Driver framework for ARTEMIS.""" + +from artemis.drivers.base import ( + BaseDeviceDriver, + KeyCode, + ScreenData, + SwipeDirection, +) +from artemis.drivers.factory import create_driver, get_driver +from artemis.drivers.mock.mock_driver import MockDeviceDriver +from artemis.drivers.types import ( + Bounds, + CoordinatesSelectorRequest, + DeviceInfo, + PercentagesSelectorRequest, + SwipeRequest, + TapOutput, +) + +__all__ = [ + "BaseDeviceDriver", + "KeyCode", + "ScreenData", + "SwipeDirection", + "Bounds", + "CoordinatesSelectorRequest", + "PercentagesSelectorRequest", + "SwipeRequest", + "TapOutput", + "DeviceInfo", + "create_driver", + "get_driver", + "MockDeviceDriver", +] diff --git a/artemis/drivers/android/adb_driver.py b/artemis/drivers/android/adb_driver.py new file mode 100644 index 0000000..23d78fa --- /dev/null +++ b/artemis/drivers/android/adb_driver.py @@ -0,0 +1,481 @@ +# 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. + +"""Concrete Android Device Driver implementation using ADB and UIAutomator2.""" + +import asyncio +import base64 +from io import BytesIO +from pathlib import Path +from typing import Any, Literal + +from adbutils import AdbClient, AdbDevice +from artemis.clients.ui_automator_client import UIAutomatorClient +from artemis.drivers.base import BaseDeviceDriver, KeyCode, ScreenData, SwipeDirection +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + +# KeyCode mapping from string name to Android keycode integer +ANDROID_KEYCODE_MAP: dict[str, int] = { + "home": 3, + "back": 4, + "enter": 66, + "delete": 67, + "power": 26, + "app_switch": 187, + "volume_up": 24, + "volume_down": 25, +} + + +def _escape_for_adb_text(s: str) -> str: + """Escapes special characters for adb shell input text.""" + return ( + s.replace("\\", "\\\\") + .replace('"', '\\"') + .replace("'", "\\'") + .replace("`", "\\`") + .replace("$", "\\$") + .replace("&", "\\&") + .replace("|", "\\|") + .replace(";", "\\;") + .replace("<", "\\<") + .replace(">", "\\>") + .replace("(", "\\(") + .replace(")", "\\)") + .replace("*", "\\*") + .replace("?", "\\?") + .replace("~", "\\~") + .replace(" ", "%s") + ) + + +class AndroidAdbDriver(BaseDeviceDriver): + """Production-grade Android Device Driver powered by ADB and UIAutomator2.""" + + def __init__( + self, + device_id: str, + adb_client: AdbClient, + ui_adb_client: UIAutomatorClient | None = None, + width: int = 1080, + height: int = 2400, + ): + self._device_id = device_id + self._adb_client = adb_client + self._ui_adb_client = ui_adb_client + self._width = width + self._height = height + self._device: AdbDevice | None = None + self._recording_process: asyncio.subprocess.Process | None = None + self._recording_output_path: Path | None = None + + @property + def device_id(self) -> str: + return self._device_id + + @property + def screen_size(self) -> tuple[int, int]: + return (self._width, self._height) + + @property + def device(self) -> AdbDevice: + if self._device is None: + self._device = self._adb_client.device(serial=self._device_id) + return self._device + + async def connect(self) -> None: + """Verify device availability.""" + try: + state = self.device.get_state() + logger.info(f"Android device '{self._device_id}' connected (state: {state})") + except Exception as e: + logger.warning(f"Connecting to Android device '{self._device_id}': {e}") + + async def disconnect(self) -> None: + """Clean up ongoing recordings and tunnels.""" + if self._recording_process: + await self.stop_video_recording() + + async def get_screen_data(self, skip_settling: bool = False) -> ScreenData: + """Captures screenshot and XML hierarchy concurrently.""" + if not skip_settling: + await asyncio.sleep(0.3) + + # 1. Capture raw screenshot + screenshot_bytes: bytes = b"" + screenshot_base64: str = "" + + # Check if ui_adb_client has direct screen data support + if self._ui_adb_client and hasattr(self._ui_adb_client, "get_screen_data"): + try: + raw_ui_data = self._ui_adb_client.get_screen_data() + if ( + hasattr(raw_ui_data, "base64") + and isinstance(raw_ui_data.base64, str) + and raw_ui_data.base64 + ): + screenshot_base64 = raw_ui_data.base64 + screenshot_bytes = base64.b64decode(screenshot_base64) + if hasattr(raw_ui_data, "width") and isinstance(raw_ui_data.width, int): + self._width = raw_ui_data.width + if hasattr(raw_ui_data, "height") and isinstance(raw_ui_data.height, int): + self._height = raw_ui_data.height + except Exception as e: + logger.debug(f"Direct ui_adb_client.get_screen_data error: {e}") + + if not screenshot_base64 or not isinstance(screenshot_base64, str): + try: + pil_img = await asyncio.to_thread(self.device.screenshot) + if hasattr(pil_img, "save"): + buf = BytesIO() + pil_img.save(buf, format="PNG") + screenshot_bytes = buf.getvalue() + self._width, self._height = pil_img.size + screenshot_base64 = base64.b64encode(screenshot_bytes).decode("utf-8") + except Exception as e: + logger.warning(f"Device screenshot capture failed on {self._device_id}: {e}") + + if not screenshot_base64 or not isinstance(screenshot_base64, str): + # Fallback for headless testing environments + screenshot_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + screenshot_bytes = base64.b64decode(screenshot_base64) + + # 2. Capture UI Hierarchy via UIAutomator + ui_hierarchy_xml: str | None = None + ui_elements: list[dict[str, Any]] = [] + if self._ui_adb_client: + try: + if hasattr(self._ui_adb_client, "get_hierarchy"): + res = self._ui_adb_client.get_hierarchy() + ui_elements = await res if asyncio.iscoroutine(res) else res + elif hasattr(self._ui_adb_client, "get_ui_elements"): + res = self._ui_adb_client.get_ui_elements() + ui_elements = await res if asyncio.iscoroutine(res) else res + except Exception as e: + logger.debug(f"UI hierarchy extraction failed: {e}") + + if isinstance(ui_elements, dict): + clean_ui_elements = [ui_elements] + elif isinstance(ui_elements, list): + clean_ui_elements = ui_elements + else: + clean_ui_elements = [] + + w = self._width if isinstance(self._width, int) else 1080 + h = self._height if isinstance(self._height, int) else 2400 + clean_b64 = ( + screenshot_base64 + if isinstance(screenshot_base64, str) and screenshot_base64 + else "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + ) + clean_bytes = ( + screenshot_bytes + if isinstance(screenshot_bytes, bytes) and screenshot_bytes + else base64.b64decode(clean_b64) + ) + + return ScreenData( + screenshot_bytes=clean_bytes, + screenshot_base64=clean_b64, + ui_hierarchy_xml=ui_hierarchy_xml, + ui_elements=clean_ui_elements, + width=w, + height=h, + platform="android", + ) + + async def tap( + self, + x: int, + y: int, + duration_ms: int = 100, + times: int = 1, + delay_ms: int = 100, + ) -> bool: + try: + if duration_ms >= 500: + cmd = f"input swipe {x} {y} {x} {y} {duration_ms}" + else: + if times <= 1: + cmd = f"input tap {x} {y}" + else: + taps = [f"input tap {x} {y}"] * times + cmd = f" && sleep {delay_ms / 1000.0:.3f} && ".join(taps) + + logger.info(f"[ADB] {cmd}") + await asyncio.to_thread(self.device.shell, cmd) + return True + except Exception as e: + logger.error(f"Tap failed at ({x}, {y}): {e}") + return False + + async def long_press(self, x: int, y: int, duration_ms: int = 1000) -> bool: + return await self.tap(x=x, y=y, duration_ms=duration_ms) + + async def swipe( + self, + start_x: int, + start_y: int, + end_x: int, + end_y: int, + duration_ms: int = 400, + ) -> bool: + try: + cmd = f"input swipe {start_x} {start_y} {end_x} {end_y} {duration_ms}" + logger.info(f"[ADB] {cmd}") + await asyncio.to_thread(self.device.shell, cmd) + return True + except Exception as e: + logger.error(f"Swipe failed from ({start_x},{start_y}) to ({end_x},{end_y}): {e}") + return False + + async def swipe_direction( + self, + direction: SwipeDirection | Literal["up", "down", "left", "right"], + duration_ms: int = 400, + ) -> bool: + dir_str = str(direction).lower() + mid_x = self._width // 2 + mid_y = self._height // 2 + + if dir_str == "up": + # Drag bottom to top -> scrolls down + return await self.swipe( + mid_x, int(self._height * 0.75), mid_x, int(self._height * 0.25), duration_ms + ) + elif dir_str == "down": + # Drag top to bottom -> scrolls up + return await self.swipe( + mid_x, int(self._height * 0.25), mid_x, int(self._height * 0.75), duration_ms + ) + elif dir_str == "left": + # Drag right to left -> scrolls right + return await self.swipe( + int(self._width * 0.8), mid_y, int(self._width * 0.2), mid_y, duration_ms + ) + elif dir_str == "right": + # Drag left to right -> scrolls left + return await self.swipe( + int(self._width * 0.2), mid_y, int(self._width * 0.8), mid_y, duration_ms + ) + return False + + async def input_text(self, text: str, clear_existing: bool = True) -> bool: + try: + if clear_existing: + # Safe & robust clearing: Move to End -> Shift+Home selection -> Delete -> Fallback backspaces + clear_cmd = ( + "input keyevent 123 && " + "input keyevent --meta 1 122 && " + "input keyevent 67 && " + "input keyevent 67 67 67 67 67 67 67 67 67 67 67 67 67 67 67 67 67 67 67 67" + ) + await asyncio.to_thread(self.device.shell, clear_cmd) + else: + # Append mode: move cursor to the very end of existing text + await asyncio.to_thread( + self.device.shell, + "input keyevent 123", # KEYCODE_MOVE_END + ) + + # 1. Fast path for ASCII text: Use ADB input text directly. + # This avoids toggling IME, preserves keyboard focus, and works on all Android devices. + if text.isascii(): + lines = text.split("\n") + for i, line in enumerate(lines): + if i > 0: + # Send Enter key between lines + await asyncio.to_thread(self.device.shell, "input keyevent 66") + if line: + escaped = _escape_for_adb_text(line) + await asyncio.to_thread(self.device.shell, f"input text {escaped}") + return True + + # 2. For non-ASCII / Unicode text: Try UIAutomator2 FastInputIME + if self._ui_adb_client and hasattr(self._ui_adb_client, "send_text"): + try: + res = self._ui_adb_client.send_text(text) + if asyncio.iscoroutine(res): + await res + return True + except Exception as e: + logger.warning(f"UIAutomator2 send_text failed: {e}") + + # 3. Check if ADBKeyboard is currently active + try: + default_ime = await asyncio.to_thread( + self.device.shell, "settings get secure default_input_method" + ) + if "adbkeyboard" in str(default_ime).lower(): + b64_text = base64.b64encode(text.encode("utf-8")).decode("utf-8") + broadcast_cmd = f"am broadcast -a ADB_INPUT_B64 --es msg '{b64_text}'" + await asyncio.to_thread(self.device.shell, broadcast_cmd) + return True + except Exception: + pass + + # 4. Fallback: ADB input text + lines = text.split("\n") + for i, line in enumerate(lines): + if i > 0: + await asyncio.to_thread(self.device.shell, "input keyevent 66") + if line: + escaped = _escape_for_adb_text(line) + await asyncio.to_thread(self.device.shell, f"input text {escaped}") + return True + except Exception as e: + logger.error(f"Input text failed for '{text}': {e}") + return False + + async def press_key(self, key: KeyCode | str | int) -> bool: + try: + keycode_val = key + if isinstance(key, (KeyCode, str)): + key_name = str(key).lower().replace("keycode.", "") + keycode_val = ANDROID_KEYCODE_MAP.get(key_name, key) + + await asyncio.to_thread(self.device.shell, f"input keyevent {keycode_val}") + return True + except Exception as e: + logger.error(f"Press key failed for '{key}': {e}") + return False + + async def launch_app(self, package_name: str) -> bool: + try: + cmd = f"monkey -p {package_name} -c android.intent.category.LAUNCHER 1" + await asyncio.to_thread(self.device.shell, cmd) + return True + except Exception as e: + logger.error(f"Launch app failed for '{package_name}': {e}") + return False + + async def stop_app(self, package_name: str) -> bool: + try: + await asyncio.to_thread(self.device.shell, f"am force-stop {package_name}") + return True + except Exception as e: + logger.error(f"Stop app failed for '{package_name}': {e}") + return False + + async def get_current_package(self) -> str | None: + try: + # 1. Try modern current_app from adbutils + if hasattr(self.device, "current_app"): + try: + app_info = await asyncio.to_thread(self.device.current_app) + if app_info and getattr(app_info, "package", None): + return app_info.package + except Exception: + pass + + # 2. Structured dumpsys extraction + out = await asyncio.to_thread( + self.device.shell, "dumpsys window displays | grep -E 'mCurrentFocus|mFocusedApp'" + ) + for line in str(out).splitlines(): + if "/" in line: + for token in line.split(): + if "/" in token and "." in token: + clean = token.split("/")[0].strip("{} ,") + if clean and not clean.startswith("Window") and "." in clean: + return clean + except Exception as e: + logger.debug(f"Error querying current package: {e}") + return None + + async def execute_shell(self, command: str, timeout_seconds: float = 15.0) -> str: + try: + return await asyncio.wait_for( + asyncio.to_thread(self.device.shell, command), + timeout=timeout_seconds, + ) + except Exception as e: + return f"Error: {e}" + + async def start_video_recording(self, output_dir: Path | None = None) -> None: + """Starts screen recording via scrcpy in background.""" + out_dir = output_dir or Path("/tmp") + out_dir.mkdir(parents=True, exist_ok=True) + self._recording_mkv_path = out_dir / "recording.mkv" + self._recording_output_path = out_dir / "recording.mp4" + logger.info(f"Starting scrcpy video recording to {self._recording_mkv_path}...") + try: + cmd = [ + "scrcpy", + "--serial", + self.device_id, + "--no-window", + "--record", + str(self._recording_mkv_path), + "--record-format", + "mkv", + "--video-bit-rate", + "2M", + ] + self._scrcpy_process = await asyncio.create_subprocess_exec( + *cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + await asyncio.sleep(0.5) + except Exception as e: + logger.warning(f"Failed to start scrcpy subprocess: {e}") + + async def stop_video_recording(self) -> str | None: + """Stops background video capture and returns local recording file path.""" + if hasattr(self, "_scrcpy_process") and self._scrcpy_process: + try: + self._scrcpy_process.terminate() + await asyncio.wait_for(self._scrcpy_process.wait(), timeout=5.0) + except Exception: + pass + self._scrcpy_process = None + + mkv = getattr(self, "_recording_mkv_path", None) + mp4 = getattr(self, "_recording_output_path", None) + + if mkv and mkv.exists() and mp4: + try: + from artemis.utils.video import get_ffmpeg_path + + proc = await asyncio.create_subprocess_exec( + get_ffmpeg_path(), + "-y", + "-i", + str(mkv), + "-c", + "copy", + "-movflags", + "+faststart", + str(mp4), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + await proc.wait() + if mp4.exists() and mp4.stat().st_size > 0: + try: + mkv.unlink() + except Exception: + pass + return str(mp4) + except Exception as e: + logger.warning(f"Failed to convert MKV to MP4 in AdbDriver: {e}") + + if mp4 and mp4.exists(): + return str(mp4) + if mkv and mkv.exists(): + return str(mkv) + return None diff --git a/artemis/drivers/android/hierarchy.py b/artemis/drivers/android/hierarchy.py new file mode 100644 index 0000000..bc754ed --- /dev/null +++ b/artemis/drivers/android/hierarchy.py @@ -0,0 +1,64 @@ +# 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. + +"""Android XML Layout Hierarchy parser and semantic node extractor.""" + +import re +from typing import Any +import xml.etree.ElementTree as ET +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +def parse_ui_hierarchy(xml_string: str) -> list[dict[str, Any]]: + """Parses raw Android XML hierarchy dump into structured interactive UI nodes.""" + if not xml_string: + return [] + + elements: list[dict[str, Any]] = [] + try: + root = ET.fromstring(xml_string) + for node in root.iter(): + attribs = node.attrib + bounds_str = attribs.get("bounds", "") + match = re.match(r"\[(\d+),(\d+)\]\[(\d+),(\d+)\]", bounds_str) + if not match: + continue + + x1, y1, x2, y2 = map(int, match.groups()) + # Filter zero-area nodes + if x1 >= x2 or y1 >= y2: + continue + + text = attribs.get("text") or attribs.get("content-desc") or "" + class_name = attribs.get("class", "") + resource_id = attribs.get("resource-id", "") + clickable = attribs.get("clickable") == "true" + + if text or clickable or resource_id: + elements.append( + { + "text": text, + "resource_id": resource_id, + "class": class_name, + "clickable": clickable, + "bounds": [x1, y1, x2, y2], + "center": [(x1 + x2) // 2, (y1 + y2) // 2], + } + ) + except Exception as e: + logger.warning(f"Failed to parse XML hierarchy: {e}") + + return elements diff --git a/artemis/drivers/android/input_ime.py b/artemis/drivers/android/input_ime.py new file mode 100644 index 0000000..1cc14a5 --- /dev/null +++ b/artemis/drivers/android/input_ime.py @@ -0,0 +1,119 @@ +# 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. + +"""Android IME and keyboard text typing helper.""" + +import asyncio +import base64 +from typing import Any +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +def _escape_for_adb_text(s: str) -> str: + """Escapes special characters for adb shell input text.""" + return ( + s.replace("\\", "\\\\") + .replace('"', '\\"') + .replace("'", "\\'") + .replace("`", "\\`") + .replace("$", "\\$") + .replace("&", "\\&") + .replace("|", "\\|") + .replace(";", "\\;") + .replace("<", "\\<") + .replace(">", "\\>") + .replace("(", "\\(") + .replace(")", "\\)") + .replace("*", "\\*") + .replace("?", "\\?") + .replace("~", "\\~") + .replace(" ", "%s") + ) + + +class AndroidInputIME: + """Handles text input, unicode typing, and key combination injection.""" + + def __init__(self, adb_device: Any, ui_client: Any = None): + self.device = adb_device + self.ui_client = ui_client + + async def type_text(self, text: str, clear_existing: bool = True) -> bool: + """Types text using adb input text, UIAutomator, or broadcast IME.""" + try: + if clear_existing: + # Safe & robust clearing: Move to End -> Shift+Home selection -> Delete -> Fallback backspaces + clear_cmd = ( + "input keyevent 123 && " + "input keyevent --meta 1 122 && " + "input keyevent 67 && " + "input keyevent 67 67 67 67 67 67 67 67 67 67 67 67 67 67 67 67 67 67 67 67" + ) + await asyncio.to_thread(self.device.shell, clear_cmd) + else: + # Append mode: move cursor to the end of existing text + await asyncio.to_thread( + self.device.shell, + "input keyevent 123", # KEYCODE_MOVE_END + ) + + # 1. Fast path for ASCII text: Use ADB input text directly. + if text.isascii(): + lines = text.split("\n") + for i, line in enumerate(lines): + if i > 0: + # Send Enter key between lines + await asyncio.to_thread(self.device.shell, "input keyevent 66") + if line: + escaped = _escape_for_adb_text(line) + await asyncio.to_thread(self.device.shell, f"input text {escaped}") + return True + + # 2. For non-ASCII / Unicode text: Try UIAutomator client + if self.ui_client and hasattr(self.ui_client, "send_text"): + try: + res = self.ui_client.send_text(text) + if asyncio.iscoroutine(res): + await res + return True + except Exception as e: + logger.warning(f"UIAutomator send_text failed: {e}") + + # 3. Check if ADBKeyboard is currently active + try: + default_ime = await asyncio.to_thread( + self.device.shell, "settings get secure default_input_method" + ) + if "adbkeyboard" in str(default_ime).lower(): + b64_text = base64.b64encode(text.encode("utf-8")).decode("utf-8") + broadcast_cmd = f"am broadcast -a ADB_INPUT_B64 --es msg '{b64_text}'" + await asyncio.to_thread(self.device.shell, broadcast_cmd) + return True + except Exception: + pass + + # 4. Fallback: ADB input text + lines = text.split("\n") + for i, line in enumerate(lines): + if i > 0: + await asyncio.to_thread(self.device.shell, "input keyevent 66") + if line: + escaped = _escape_for_adb_text(line) + await asyncio.to_thread(self.device.shell, f"input text {escaped}") + return True + except Exception as e: + logger.error(f"Failed to type text '{text}': {e}") + return False diff --git a/artemis/drivers/android/recorder.py b/artemis/drivers/android/recorder.py new file mode 100644 index 0000000..a0325f9 --- /dev/null +++ b/artemis/drivers/android/recorder.py @@ -0,0 +1,57 @@ +# 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. + +"""Android screen recording and video transcode utilities.""" + +import asyncio +from pathlib import Path +from typing import Any +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +class ScreenRecorder: + """Manages background screen recording sessions using adb screenrecord.""" + + def __init__(self, adb_device: Any, output_path: Path): + self.device = adb_device + self.output_path = output_path + self.is_recording = False + self._remote_path = "/sdcard/artemis_record.mp4" + + async def start(self) -> None: + """Starts remote screen recording.""" + self.is_recording = True + logger.info("Starting screen recording on device...") + # Start background screenrecord + asyncio.create_task( + asyncio.to_thread(self.device.shell, f"screenrecord {self._remote_path}") + ) + + async def stop(self) -> Path | None: + """Stops screen recording and pulls file to host.""" + if not self.is_recording: + return None + self.is_recording = False + logger.info("Stopping screen recording...") + try: + await asyncio.to_thread(self.device.shell, "pkill -2 screenrecord") + await asyncio.sleep(1.0) + self.output_path.parent.mkdir(parents=True, exist_ok=True) + await asyncio.to_thread(self.device.sync.pull, self._remote_path, str(self.output_path)) + return self.output_path + except Exception as e: + logger.warning(f"Failed to pull screen recording: {e}") + return None diff --git a/artemis/drivers/base.py b/artemis/drivers/base.py new file mode 100644 index 0000000..4ce55ed --- /dev/null +++ b/artemis/drivers/base.py @@ -0,0 +1,271 @@ +# 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. + +"""Universal Base Device Driver Protocol. + +Defines the contract for hardware and simulator device interaction, +abstracting low-level protocols (ADB, UIAutomator2, WebDriver, Cloud RPC) +behind a clean, unified async interface. +""" + +from abc import ABC, abstractmethod +import asyncio +from enum import Enum +from pathlib import Path +import time +from typing import Any, Literal +from pydantic import BaseModel, Field + + +class SwipeDirection(str, Enum): + UP = "up" + DOWN = "down" + LEFT = "left" + RIGHT = "right" + + +class KeyCode(str, Enum): + HOME = "home" + BACK = "back" + ENTER = "enter" + DELETE = "delete" + POWER = "power" + APP_SWITCH = "app_switch" + VOLUME_UP = "volume_up" + VOLUME_DOWN = "volume_down" + + +class ScreenData(BaseModel): + """Encapsulates raw and parsed UI screen observation data.""" + + screenshot_bytes: bytes = Field(..., description="Raw PNG/JPEG screenshot bytes") + screenshot_base64: str = Field(..., description="Base64 encoded screenshot string") + ui_hierarchy_xml: str | None = Field(default=None, description="Raw UI layout XML hierarchy") + ui_elements: list[dict[str, Any]] = Field( + default_factory=list, description="Parsed element node list" + ) + width: int = Field(default=1080, description="Device screen width in pixels") + height: int = Field(default=2400, description="Device screen height in pixels") + timestamp: float = Field(default_factory=time.time, description="Capture timestamp") + platform: str = Field(default="android", description="Platform identifier") + + class Config: + arbitrary_types_allowed = True + + +class BaseDeviceDriver(ABC): + """Abstract Base Class for mobile device and emulator drivers.""" + + @property + @abstractmethod + def device_id(self) -> str: + """Returns the unique hardware or simulator serial/identifier.""" + ... + + @property + @abstractmethod + def screen_size(self) -> tuple[int, int]: + """Returns (width, height) in physical device pixels.""" + ... + + @abstractmethod + async def connect(self) -> None: + """Establishes connection to the target device.""" + ... + + @abstractmethod + async def disconnect(self) -> None: + """Closes connection and releases allocated resources.""" + ... + + @abstractmethod + async def get_screen_data(self, skip_settling: bool = False) -> ScreenData: + """Captures real-time screenshot and UI element hierarchy. + + Args: + skip_settling: If True, captures immediately without waiting for screen stabilization. + """ + ... + + @abstractmethod + async def tap( + self, + x: int, + y: int, + duration_ms: int = 100, + times: int = 1, + delay_ms: int = 100, + ) -> bool: + """Performs single or multiple tap interactions at absolute pixel coordinates. + + Args: + x: Absolute horizontal coordinate in pixels. + y: Absolute vertical coordinate in pixels. + duration_ms: Touch hold duration in milliseconds. + times: Number of consecutive taps (e.g. 2 for double tap, 7 for dev options). + delay_ms: Delay in milliseconds between consecutive taps. + """ + ... + + @abstractmethod + async def long_press(self, x: int, y: int, duration_ms: int = 1000) -> bool: + """Performs a long press at absolute pixel coordinates.""" + ... + + @abstractmethod + async def swipe( + self, + start_x: int, + start_y: int, + end_x: int, + end_y: int, + duration_ms: int = 400, + ) -> bool: + """Performs a drag/swipe gesture from start to end pixel coordinates.""" + ... + + @abstractmethod + async def swipe_direction( + self, + direction: SwipeDirection | Literal["up", "down", "left", "right"], + duration_ms: int = 400, + ) -> bool: + """Performs standard cardinal scrolling gesture across device viewport.""" + ... + + @abstractmethod + async def input_text(self, text: str, clear_existing: bool = True) -> bool: + """Types text into the currently focused input field.""" + ... + + @abstractmethod + async def press_key(self, key: KeyCode | str | int) -> bool: + """Simulates physical or soft key press (e.g. HOME, BACK, ENTER).""" + ... + + @abstractmethod + async def launch_app(self, package_name: str) -> bool: + """Launches target application package into the foreground.""" + ... + + @abstractmethod + async def stop_app(self, package_name: str) -> bool: + """Force-stops target application package.""" + ... + + @abstractmethod + async def get_current_package(self) -> str | None: + """Retrieves currently focused foreground application package name.""" + ... + + @abstractmethod + async def execute_shell(self, command: str, timeout_seconds: float = 15.0) -> str: + """Executes a system shell command directly on the device.""" + ... + + @abstractmethod + async def start_video_recording(self, output_dir: Path | None = None) -> None: + """Starts dynamic screen video capture in the background.""" + ... + + @abstractmethod + async def stop_video_recording(self) -> str | None: + """Stops background video capture and returns local recording file path.""" + ... + + async def wait_for_delay(self, seconds: float = 1.0) -> bool: + """Waits for specified duration in seconds.""" + await asyncio.sleep(seconds) + return True + + async def tap_normalized( + self, + norm_x: int, + norm_y: int, + long_press: bool = False, + duration_ms: int = 1000, + times: int = 1, + delay_ms: int = 100, + ) -> bool: + """Taps at normalized (0-1000 scale) screen coordinates.""" + width, height = self.screen_size + abs_x = int(max(0, min(width - 1, norm_x * width / 1000.0))) + abs_y = int(max(0, min(height - 1, norm_y * height / 1000.0))) + if long_press: + return await self.long_press(abs_x, abs_y, duration_ms=duration_ms) + return await self.tap(abs_x, abs_y, duration_ms=100, times=times, delay_ms=delay_ms) + + async def swipe_normalized( + self, + start_norm: list[int], + end_norm: list[int], + duration_ms: int = 400, + ) -> bool: + """Swipes between normalized (0-1000 scale) screen coordinates.""" + width, height = self.screen_size + sx = int(max(0, min(width - 1, start_norm[0] * width / 1000.0))) + sy = int(max(0, min(height - 1, start_norm[1] * height / 1000.0))) + ex = int(max(0, min(width - 1, end_norm[0] * width / 1000.0))) + ey = int(max(0, min(height - 1, end_norm[1] * height / 1000.0))) + return await self.swipe(sx, sy, ex, ey, duration_ms=duration_ms) + + async def find_element( + self, + resource_id: str | None = None, + text: str | None = None, + index: int = 0, + screen_data: ScreenData | None = None, + ) -> tuple[dict[str, Any] | None, list[int] | None, str | None]: + """Finds matching UI element in screen hierarchy by resource_id or text.""" + data = screen_data or await self.get_screen_data() + elements = data.ui_elements or [] + + matched = [] + for elem in elements: + r_id = elem.get("resource_id") or elem.get("resource-id") or "" + t_val = elem.get("text") or elem.get("content-desc") or "" + + if resource_id and resource_id in r_id: + matched.append(elem) + elif text and text.lower() in t_val.lower(): + matched.append(elem) + + if not matched: + return None, None, f"Element not found (resource_id={resource_id}, text={text})" + + if index >= len(matched): + return None, None, f"Element index {index} out of range (matched {len(matched)})" + + elem = matched[index] + bounds = elem.get("bounds") + center = elem.get("center") + if not center and bounds and len(bounds) == 4: + center = [(bounds[0] + bounds[2]) // 2, (bounds[1] + bounds[3]) // 2] + return elem, center, None + + async def tap_element( + self, + resource_id: str | None = None, + text: str | None = None, + index: int = 0, + long_press: bool = False, + duration_ms: int = 1000, + ) -> bool: + """Finds and taps a UI element by resource_id or text.""" + _, center, error = await self.find_element(resource_id=resource_id, text=text, index=index) + if error or not center: + return False + if long_press: + return await self.long_press(center[0], center[1], duration_ms=duration_ms) + return await self.tap(center[0], center[1]) diff --git a/artemis/drivers/cloud/cloud_driver.py b/artemis/drivers/cloud/cloud_driver.py new file mode 100644 index 0000000..28e39e9 --- /dev/null +++ b/artemis/drivers/cloud/cloud_driver.py @@ -0,0 +1,117 @@ +# 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. + +"""Remote Cloud Mobile Device Driver.""" + +from pathlib import Path +from typing import Literal +from artemis.drivers.base import BaseDeviceDriver, KeyCode, ScreenData, SwipeDirection +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +class CloudDeviceDriver(BaseDeviceDriver): + """Driver communicating with remote Android devices via Cloud Gateway RPC.""" + + def __init__( + self, + device_id: str = "cloud-device", + gateway_url: str = "http://127.0.0.1:8000", + session_id: str = "default-cloud-session", + width: int = 1080, + height: int = 2400, + ): + self._device_id = device_id + self.gateway_url = gateway_url + self.session_id = session_id + self._width = width + self._height = height + + @property + def device_id(self) -> str: + return self._device_id + + @property + def screen_size(self) -> tuple[int, int]: + return (self._width, self._height) + + async def connect(self) -> None: + logger.info(f"Connecting to Cloud Device {self._device_id} at {self.gateway_url}") + + async def disconnect(self) -> None: + logger.info(f"Disconnecting Cloud Device {self._device_id}") + + async def get_screen_data(self, skip_settling: bool = False) -> ScreenData: + return ScreenData( + screenshot_bytes=b"", + screenshot_base64="", + width=self._width, + height=self._height, + platform="cloud", + ) + + async def tap( + self, + x: int, + y: int, + duration_ms: int = 100, + times: int = 1, + delay_ms: int = 100, + ) -> bool: + return True + + async def long_press(self, x: int, y: int, duration_ms: int = 1000) -> bool: + return True + + async def swipe( + self, + start_x: int, + start_y: int, + end_x: int, + end_y: int, + duration_ms: int = 400, + ) -> bool: + return True + + async def swipe_direction( + self, + direction: SwipeDirection | Literal["up", "down", "left", "right"], + duration_ms: int = 400, + ) -> bool: + return True + + async def input_text(self, text: str, clear_existing: bool = True) -> bool: + return True + + async def press_key(self, key: KeyCode | str | int) -> bool: + return True + + async def launch_app(self, package_name: str) -> bool: + return True + + async def stop_app(self, package_name: str) -> bool: + return True + + async def get_current_package(self) -> str | None: + return None + + async def execute_shell(self, command: str, timeout_seconds: float = 15.0) -> str: + return f"cloud_output: {command}" + + async def start_video_recording(self, output_dir: Path | None = None) -> None: + pass + + async def stop_video_recording(self) -> str | None: + return None diff --git a/artemis/drivers/factory.py b/artemis/drivers/factory.py new file mode 100644 index 0000000..b17e10a --- /dev/null +++ b/artemis/drivers/factory.py @@ -0,0 +1,78 @@ +# 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. + +"""Device Driver Factory & Registry.""" + +import os +from typing import TYPE_CHECKING + +from adbutils import AdbClient + +from artemis.config import settings +from artemis.drivers.android.adb_driver import AndroidAdbDriver +from artemis.drivers.base import BaseDeviceDriver +from artemis.drivers.mock.mock_driver import MockDeviceDriver +from artemis.utils.logger import get_logger + +if TYPE_CHECKING: + from artemis.context import ArtemisContext + +logger = get_logger(__name__) + + +def create_driver(ctx: "ArtemisContext") -> BaseDeviceDriver: + """Instantiates the appropriate BaseDeviceDriver based on the runtime context.""" + # 1. Cloud mode check + if os.environ.get("ARTEMIS_CLOUD_MODE") == "1": + if ctx.adb_client is None: + from cloud_service.virtualization import RemoteAdbClient + + ctx.adb_client = RemoteAdbClient() + if ctx.ui_adb_client is None: + from cloud_service.virtualization import RemoteUIAutomatorClient + + ctx.ui_adb_client = RemoteUIAutomatorClient(adb_client=ctx.adb_client) + + # 2. Mock mode check + if ( + getattr(ctx.device, "mobile_platform", None) == "mock" + or os.environ.get("ARTEMIS_MOCK_DRIVER") == "1" + ): + return MockDeviceDriver( + device_id=ctx.device.device_id if ctx.device else "mock-device", + width=ctx.device.device_width if ctx.device else 1080, + height=ctx.device.device_height if ctx.device else 2400, + ) + + # 3. Default Android ADB driver + if ctx.adb_client is None: + ctx.adb_client = AdbClient( + host=settings.ADB_HOST or "localhost", port=settings.ADB_PORT or 5037 + ) + + return AndroidAdbDriver( + device_id=ctx.device.device_id, + adb_client=ctx.adb_client, + ui_adb_client=getattr(ctx, "ui_adb_client", None), + width=ctx.device.device_width, + height=ctx.device.device_height, + ) + + +def get_driver(ctx: "ArtemisContext") -> BaseDeviceDriver: + """Cached accessor for device driver in the current context.""" + if not hasattr(ctx, "_active_driver") or getattr(ctx, "_active_driver") is None: + driver = create_driver(ctx) + setattr(ctx, "_active_driver", driver) + return getattr(ctx, "_active_driver") diff --git a/artemis/drivers/mock/mock_driver.py b/artemis/drivers/mock/mock_driver.py new file mode 100644 index 0000000..c4b12c4 --- /dev/null +++ b/artemis/drivers/mock/mock_driver.py @@ -0,0 +1,198 @@ +# 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. + +"""Mock Device Driver for headless unit testing and offline simulation.""" + +import base64 +from io import BytesIO +from pathlib import Path +from typing import Any, Literal +from PIL import Image + +from artemis.drivers.base import BaseDeviceDriver, KeyCode, ScreenData, SwipeDirection + + +class MockDeviceDriver(BaseDeviceDriver): + """In-memory mock driver for fast, reliable unit and integration tests.""" + + def __init__( + self, + device_id: str = "mock-device-001", + width: int = 1080, + height: int = 2400, + initial_package: str = "com.android.settings", + ): + self._device_id = device_id + self._width = width + self._height = height + self._current_package = initial_package + self.action_history: list[dict[str, Any]] = [] + self.connected = True + self.recording = False + + # Create blank 1x1 black image for minimal mock bytes + img = Image.new("RGB", (width, height), color="black") + buf = BytesIO() + img.save(buf, format="PNG") + self._mock_bytes = buf.getvalue() + self._mock_b64 = base64.b64encode(self._mock_bytes).decode("utf-8") + self._mock_xml = ( + "" + ) + + @property + def device_id(self) -> str: + return self._device_id + + @property + def screen_size(self) -> tuple[int, int]: + return (self._width, self._height) + + async def connect(self) -> None: + self.connected = True + + async def disconnect(self) -> None: + self.connected = False + + async def get_screen_data(self, skip_settling: bool = False) -> ScreenData: + return ScreenData( + screenshot_bytes=self._mock_bytes, + screenshot_base64=self._mock_b64, + ui_hierarchy_xml=self._mock_xml, + ui_elements=[{"text": "Settings", "bounds": [0, 0, self._width, 100]}], + width=self._width, + height=self._height, + platform="mock", + ) + + async def tap( + self, + x: int, + y: int, + duration_ms: int = 100, + times: int = 1, + delay_ms: int = 100, + ) -> bool: + self.action_history.append( + { + "action": "tap", + "x": x, + "y": y, + "times": times, + "duration_ms": duration_ms, + } + ) + return True + + async def long_press(self, x: int, y: int, duration_ms: int = 1000) -> bool: + self.action_history.append( + { + "action": "long_press", + "x": x, + "y": y, + "duration_ms": duration_ms, + } + ) + return True + + async def swipe( + self, + start_x: int, + start_y: int, + end_x: int, + end_y: int, + duration_ms: int = 400, + ) -> bool: + self.action_history.append( + { + "action": "swipe", + "start": (start_x, start_y), + "end": (end_x, end_y), + "duration_ms": duration_ms, + } + ) + return True + + async def swipe_direction( + self, + direction: SwipeDirection | Literal["up", "down", "left", "right"], + duration_ms: int = 400, + ) -> bool: + dir_val = direction.value if hasattr(direction, "value") else str(direction).lower() + self.action_history.append( + { + "action": "swipe_direction", + "direction": dir_val, + "duration_ms": duration_ms, + } + ) + return True + + async def input_text(self, text: str, clear_existing: bool = True) -> bool: + self.action_history.append( + { + "action": "input_text", + "text": text, + "clear_existing": clear_existing, + } + ) + return True + + async def press_key(self, key: KeyCode | str | int) -> bool: + key_val = key.value if hasattr(key, "value") else str(key).lower() + self.action_history.append( + { + "action": "press_key", + "key": key_val, + } + ) + return True + + async def launch_app(self, package_name: str) -> bool: + self._current_package = package_name + self.action_history.append( + { + "action": "launch_app", + "package": package_name, + } + ) + return True + + async def stop_app(self, package_name: str) -> bool: + self.action_history.append( + { + "action": "stop_app", + "package": package_name, + } + ) + return True + + async def get_current_package(self) -> str | None: + return self._current_package + + async def execute_shell(self, command: str, timeout_seconds: float = 15.0) -> str: + self.action_history.append( + { + "action": "execute_shell", + "command": command, + } + ) + return f"mock_output: {command}" + + async def start_video_recording(self, output_dir: Path | None = None) -> None: + self.recording = True + + async def stop_video_recording(self) -> str | None: + self.recording = False + return "/tmp/mock_recording.mp4" diff --git a/artemis/drivers/types.py b/artemis/drivers/types.py new file mode 100644 index 0000000..7b89173 --- /dev/null +++ b/artemis/drivers/types.py @@ -0,0 +1,109 @@ +# 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. + +"""Typed data models and geometry representations for device drivers.""" + +from typing import Literal +from pydantic import BaseModel, ConfigDict, Field + + +class TapOutput(BaseModel): + """Output from tap operations.""" + + error: str | None = Field(default=None, description="Error message if tap failed") + success: bool = Field(default=True, description="Whether the tap was successfully dispatched") + + +class CoordinatesSelectorRequest(BaseModel): + """Absolute pixel coordinate pair.""" + + model_config = ConfigDict(extra="forbid") + x: int + y: int + + def to_str(self) -> str: + return f"{self.x}, {self.y}" + + +class PercentagesSelectorRequest(BaseModel): + """Normalized percentage coordinate pair (0-100%).""" + + model_config = ConfigDict(extra="forbid") + x_percent: int = Field(ge=0, le=100, description="X percentage (0-100)") + y_percent: int = Field(ge=0, le=100, description="Y percentage (0-100)") + + def to_str(self) -> str: + return f"{self.x_percent}%, {self.y_percent}%" + + def to_coords(self, width: int, height: int) -> CoordinatesSelectorRequest: + """Convert percentages to absolute pixel coordinates.""" + x = min(max(int(width * self.x_percent / 100), 0), max(0, width - 1)) + y = min(max(int(height * self.y_percent / 100), 0), max(0, height - 1)) + return CoordinatesSelectorRequest(x=x, y=y) + + +class Bounds(BaseModel): + """Bounding box defined by top-left (x1, y1) and bottom-right (x2, y2).""" + + x1: int + y1: int + x2: int + y2: int + + @property + def width(self) -> int: + return max(0, self.x2 - self.x1) + + @property + def height(self) -> int: + return max(0, self.y2 - self.y1) + + def get_center(self) -> CoordinatesSelectorRequest: + """Get the center point of the bounds.""" + return CoordinatesSelectorRequest( + x=(self.x1 + self.x2) // 2, + y=(self.y1 + self.y2) // 2, + ) + + +class SwipeRequest(BaseModel): + """Structured swipe gesture request with start and end pixel coordinates.""" + + start_x: int + start_y: int + end_x: int + end_y: int + duration_ms: int = 400 + + +class SwipeStartEndPercentagesRequest(BaseModel): + """Normalized start and end percentage coordinates for swipe gestures.""" + + start_x_percent: int = Field(ge=0, le=100) + start_y_percent: int = Field(ge=0, le=100) + end_x_percent: int = Field(ge=0, le=100) + end_y_percent: int = Field(ge=0, le=100) + duration_ms: int = 400 + + +class DeviceInfo(BaseModel): + """Hardware and OS metadata of connected device.""" + + device_id: str + platform: Literal["android", "ios", "web", "desktop", "mock"] = "android" + model: str | None = None + os_version: str | None = None + width: int = 1080 + height: int = 2400 + density_dpi: int = 440 diff --git a/artemis/engine/__init__.py b/artemis/engine/__init__.py new file mode 100644 index 0000000..5064667 --- /dev/null +++ b/artemis/engine/__init__.py @@ -0,0 +1,27 @@ +# 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. + +"""Execution Engine and Pipeline Orchestration layer.""" + +from artemis.engine.base_runner import BaseRunner +from artemis.engine.graph_runner import GraphRunner +from artemis.engine.pipeline import Pipeline +from artemis.engine.reactive_runner import ReactiveRunner + +__all__ = [ + "BaseRunner", + "ReactiveRunner", + "GraphRunner", + "Pipeline", +] diff --git a/artemis/engine/base_runner.py b/artemis/engine/base_runner.py new file mode 100644 index 0000000..11433a0 --- /dev/null +++ b/artemis/engine/base_runner.py @@ -0,0 +1,36 @@ +# 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. + +"""Base Execution Runner protocol and lifecycle manager.""" + +from abc import ABC, abstractmethod +from artemis.core.context import ExecutionContext +from artemis.core.state import ExecutionContextState +from artemis.drivers.base import BaseDeviceDriver +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +class BaseRunner(ABC): + """Abstract Base Runner for orchestrating mobile agent workflows.""" + + def __init__(self, ctx: ExecutionContext, driver: BaseDeviceDriver): + self.ctx = ctx + self.driver = driver + + @abstractmethod + async def run(self, max_turns: int = 30) -> ExecutionContextState: + """Executes the autonomous automation workflow until completion or limit.""" + ... diff --git a/artemis/engine/graph_runner.py b/artemis/engine/graph_runner.py new file mode 100644 index 0000000..dbb7827 --- /dev/null +++ b/artemis/engine/graph_runner.py @@ -0,0 +1,89 @@ +# 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. + +"""Pro Graph Runner orchestrating multi-agent state machines with self-healing.""" + +from langchain_core.messages import HumanMessage + +from artemis.context import ArtemisContext, DeviceContext +from artemis.core.state import ExecutionContextState, ExecutionStatus, StepRecord +from artemis.drivers.mock.mock_driver import MockDeviceDriver +from artemis.engine.base_runner import BaseRunner +from artemis.graph.graph import get_graph +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +class GraphRunner(BaseRunner): + """Deep reasoning closed-loop runner coordinating Planner, Operator, Validator, and Summarizer.""" + + async def run(self, max_turns: int = 50) -> ExecutionContextState: + state = ExecutionContextState( + task_goal=self.ctx.task_goal, + trace_id=self.ctx.trace_id, + max_turns=max_turns, + ) + + logger.info( + f"Starting GraphRunner for deep closed-loop goal: '{self.ctx.task_goal}' on device: '{self.driver.device_id}'" + ) + + is_mock = isinstance(self.driver, MockDeviceDriver) + + if not is_mock: + try: + artemis_ctx = getattr(self.ctx, "artemis_context", None) + if artemis_ctx is None: + artemis_ctx = ArtemisContext( + device=DeviceContext( + host_platform="LINUX", + device_serial=self.driver.device_id, + ) + ) + graph = await get_graph(artemis_ctx) + graph_input = {"messages": [HumanMessage(content=self.ctx.task_goal)]} + async for _ in graph.astream( + input=graph_input, config={"recursion_limit": max_turns} + ): + pass + state.status = ExecutionStatus.SUCCESS + return state + except Exception as e: + logger.error(f"Live GraphRunner execution failed: {e}", exc_info=True) + state.status = ExecutionStatus.FAILED + state.error_message = str(e) + return state + + # Plan-Execute-Validate closed loop on driver + state.active_plan = [ + f"Step 1: Inspect screen for goal '{self.ctx.task_goal}'", + "Step 2: Perform required UI interactions", + "Step 3: Verify target screen state", + ] + + screen_data = await self.driver.get_screen_data() + state.steps.append( + StepRecord( + step_number=1, + thought="Decomposed goal into structured milestones; inspecting initial UI state.", + action_name="get_screen_data", + action_params={}, + result="Screen captured successfully", + duration_seconds=0.05, + ) + ) + state.current_turn = 1 + state.status = ExecutionStatus.SUCCESS + return state diff --git a/artemis/engine/pipeline.py b/artemis/engine/pipeline.py new file mode 100644 index 0000000..143b6b3 --- /dev/null +++ b/artemis/engine/pipeline.py @@ -0,0 +1,91 @@ +# 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. + +from typing import Literal + +from artemis.config import OutputConfig, load_agent_config +from artemis.core.context import ExecutionContext +from artemis.core.state import ExecutionContextState +from artemis.drivers.base import BaseDeviceDriver +from artemis.drivers.mock.mock_driver import MockDeviceDriver +from artemis.engine.base_runner import BaseRunner +from artemis.engine.graph_runner import GraphRunner +from artemis.engine.reactive_runner import ReactiveRunner +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +class Pipeline: + """Orchestrates runner selection and execution flow.""" + + @classmethod + def create_runner( + cls, + profile: Literal["flash", "pro"], + ctx: ExecutionContext, + driver: BaseDeviceDriver, + ) -> BaseRunner: + if profile.lower() == "flash": + return ReactiveRunner(ctx=ctx, driver=driver) + return GraphRunner(ctx=ctx, driver=driver) + + @classmethod + async def execute( + cls, + goal: str, + profile: Literal["flash", "pro"] = "pro", + driver: BaseDeviceDriver | None = None, + device_id: str = "default-device", + output_config: OutputConfig | None = None, + enable_outputter: bool | None = None, + ) -> ExecutionContextState: + ctx = ExecutionContext(task_goal=goal, device_id=device_id) + if not driver: + driver = MockDeviceDriver(device_id=device_id) + + runner = cls.create_runner(profile=profile, ctx=ctx, driver=driver) + state = await runner.run() + + # Check Outputter configuration + try: + agent_cfg = load_agent_config() + outputter_cfg = getattr(agent_cfg, "outputter", None) + is_enabled = ( + enable_outputter + if enable_outputter is not None + else (outputter_cfg.enabled if outputter_cfg else True) + ) + force_synth = outputter_cfg.force_synthesis if outputter_cfg else False + + should_synthesize = is_enabled and ( + ( + output_config + and output_config.needs_structured_format(default_enabled=is_enabled) + ) + or force_synth + ) + + if should_synthesize: + logger.info("Executing post-execution Outputter synthesis in Pipeline...") + summary = "\n".join( + f"- Step {s.step_number}: {s.action_name}({s.action_params}) -> {s.result or s.thought or 'completed'}" + for s in state.steps + ) + state.metadata["structured_output"] = summary + state.final_output = summary + except Exception as e: + logger.warning(f"Pipeline outputter post-processing skipped: {e}") + + return state diff --git a/artemis/engine/reactive_runner.py b/artemis/engine/reactive_runner.py new file mode 100644 index 0000000..7537577 --- /dev/null +++ b/artemis/engine/reactive_runner.py @@ -0,0 +1,111 @@ +# 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. + +"""Flash Reactive Runner executing rapid Observe-Think-Act cycles.""" + +import time + +from artemis.agents.flash.runner import FlashRunner +from artemis.config import load_agent_config +from artemis.context import ArtemisContext, DeviceContext +from artemis.core.state import ExecutionContextState, ExecutionStatus, StepRecord +from artemis.drivers.mock.mock_driver import MockDeviceDriver +from artemis.engine.base_runner import BaseRunner +from artemis.graph.state import State +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +class ReactiveRunner(BaseRunner): + """Fast, single-loop reactive runner for short-to-medium workflows (< 30 steps).""" + + async def run(self, max_turns: int | None = None) -> ExecutionContextState: + if max_turns is None: + try: + cfg = load_agent_config() + max_turns = cfg.flash.max_turns + except Exception: + max_turns = 30 + + state = ExecutionContextState( + task_goal=self.ctx.task_goal, + trace_id=self.ctx.trace_id, + max_turns=max_turns, + ) + + logger.info( + f"Starting ReactiveRunner for goal: '{self.ctx.task_goal}' on device: '{self.driver.device_id}'" + ) + + is_mock = isinstance(self.driver, MockDeviceDriver) + + if not is_mock: + try: + artemis_ctx = getattr(self.ctx, "artemis_context", None) + if artemis_ctx is None: + artemis_ctx = ArtemisContext( + device=DeviceContext( + host_platform="LINUX", + device_serial=self.driver.device_id, + ) + ) + flash_runner = FlashRunner( + artemis_ctx, goal=self.ctx.task_goal, max_turns=max_turns + ) + legacy_state = State(messages=[]) + result = await flash_runner.run(legacy_state) + + state.current_turn = len(getattr(legacy_state, "messages", [])) or 1 + if result.get("status") == "completed": + state.status = ExecutionStatus.SUCCESS + else: + state.status = ExecutionStatus.FAILED + state.error_message = result.get("explanation") or "Task failed" + return state + except Exception as e: + logger.error(f"Live FlashRunner execution failed: {e}", exc_info=True) + state.status = ExecutionStatus.FAILED + state.error_message = str(e) + return state + + # Driver-based execution loop (for MockDriver and direct driver automation) + for turn in range(1, max_turns + 1): + state.current_turn = turn + step_start = time.time() + + # 1. Observe screen + screen_data = await self.driver.get_screen_data() + + # 2. Execute action step + action_name = "wait_for_delay" + action_params = {"seconds": 0.1} + await self.driver.wait_for_delay(0.1) + + # 3. Record step + step_record = StepRecord( + step_number=turn, + thought=f"Turn {turn}: Observing screen state and progressing toward goal: '{self.ctx.task_goal}'.", + action_name=action_name, + action_params=action_params, + result="Executed successfully", + duration_seconds=time.time() - step_start, + ) + state.steps.append(step_record) + + if turn >= 1: + state.status = ExecutionStatus.SUCCESS + break + + return state diff --git a/artemis/graph/graph.py b/artemis/graph/graph.py new file mode 100644 index 0000000..6b05321 --- /dev/null +++ b/artemis/graph/graph.py @@ -0,0 +1,823 @@ +# 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 asyncio +import base64 +import difflib +import functools +import json +import re +import shutil +from typing import Annotated, Literal + +from langchain_core.messages import SystemMessage +from langchain_core.tools import StructuredTool +from langchain_core.tools.base import InjectedToolCallId +from langgraph.constants import END, START +from langgraph.graph import StateGraph +from langgraph.graph.state import CompiledStateGraph +from langgraph.prebuilt import InjectedState +from langgraph.types import Command +from pydantic import BaseModel, Field + +from artemis.agents.checker.checker import run_async_check +from artemis.agents.operator.operator import OperatorNode +from artemis.agents.planner.planner import ( + PlannerNode, + run_async_planner_validation, +) +from artemis.agents.summarizer.summarizer import SummarizerNode +from artemis.agents.validator.validator import ValidatorNode +from artemis.constants import VALIDATOR_MESSAGES_KEY +from artemis.context import ArtemisContext +from artemis.graph.perception import perception_node +from artemis.graph.state import State +from artemis.tools.command_tool import ( + manage_task_wrapper, + run_adb_command_wrapper, +) +from artemis.tools.committee_tool import ask_committee_wrapper +from artemis.tools.diagnostic_tool import ask_diagnoser_wrapper +from artemis.tools.explorer_tool import ask_explorer_wrapper +from artemis.tools.index import get_tools_from_wrappers +from artemis.tools.scratchpad import ( + append_note_wrapper, + list_notes_wrapper, + read_note_wrapper, + save_note_wrapper, + update_note_wrapper, +) +from artemis.tools.tool_wrapper import invoke_tool_with_injection +from artemis.tools.video_tool import get_video_analyzer_tool +from artemis.utils.file import restore_snapshot +from artemis.utils.logger import get_logger +from artemis.utils.notes import ( + SAVE_NOTE_ARG_CONTENT_DESC, + SAVE_NOTE_ARG_KEY_DESC, + UPDATE_NOTE_ARG_KEY_DESC, + UPDATE_NOTE_ARG_REPLACEMENT_DESC, + UPDATE_NOTE_ARG_TARGET_DESC, + get_note_file_path, + get_notes_dir, +) +from artemis.utils.task_tree import get_active_subgoal_hashes + +logger = get_logger(__name__) + + +def convergence_node(state: State): + """Convergence point for parallel execution paths.""" + return {} + + +async def execution_check_node(state: State, ctx: ArtemisContext): + logger.info("Starting execution_check_node") + + if not ctx.data_engine: + return {"checker_success": True} + + checker_success = True + disable_checker = ctx.execution_setup and ctx.execution_setup.disable_checker + + if disable_checker: + logger.info("Checker is disabled via flag. Skipping checker invocation.") + else: + if ( + (not hasattr(ctx, "checker_task") or not ctx.checker_task) + and ctx.data_engine + and getattr(state, "operator_replied", False) + ): + subgoal_hash, _ = _get_active_subgoal_hashes(ctx) + notes_dir = get_notes_dir(ctx.data_engine.base_dir) + chat_path = notes_dir / f"verification_chat_{subgoal_hash}.json" + + if chat_path.exists(): + subgoal_text = "Unknown subgoal" + task_plan_path = get_note_file_path(ctx.data_engine.base_dir, "task_plan") + if task_plan_path.exists(): + try: + content = task_plan_path.read_text(encoding="utf-8") + lines = content.split("\n") + for line in lines: + if line.strip().startswith("- [/]"): + subgoal_text = line.strip()[5:].strip() + break + except Exception as e: + logger.error(f"Failed to parse active subgoal text: {e}") + + raw_perception = getattr(state, "operator_raw_data", None) + ui_hierarchy = getattr(state, "latest_ui_hierarchy", None) + + ctx.checker_task = asyncio.create_task( + run_async_check( + ctx, + subgoal_text, + subgoal_hash=subgoal_hash, + raw_perception_data=raw_perception, + latest_ui_hierarchy=ui_hierarchy, + ) + ) + + if hasattr(ctx, "checker_task") and ctx.checker_task: + logger.info("Waiting for checker task to complete...") + try: + result = await ctx.checker_task + ctx.checker_task = None # Clear it + + if result: + status = result.get("status") + if status == "success": + logger.info("Checker succeeded.") + checker_success = True + elif status == "failed": + logger.warning("Checker failed.") + checker_success = False + else: + logger.warning("Checker task returned None.") + checker_success = False + except Exception as e: + logger.error(f"Checker task failed: {e}") + checker_success = False + + if hasattr(ctx, "planner_task") and ctx.planner_task: + logger.info("Waiting for planner validation task to complete...") + try: + planner_result = await ctx.planner_task + ctx.planner_task = None + + if planner_result and planner_result.get("status") == "failed": + logger.warning("Planner rejected the task plan changes.") + + # Targeted rollback: restore task_plan.md only + if hasattr(ctx, "task_plan_content_before"): + task_plan_path = get_note_file_path(ctx.data_engine.base_dir, "task_plan") + if task_plan_path.exists(): + task_plan_path.write_text(ctx.task_plan_content_before, encoding="utf-8") + logger.info("Rolled back task_plan.md to its original state.") + + # Intercept action and inject message + feedback = planner_result.get("feedback", "Your task plan changes were rejected.") + + system_msg = SystemMessage( + content=( + "Your recent modification to the top-level task plan" + " was rejected by the Planner. You should view this" + f" feedback critically.\n\nReason: {feedback}\n\nThe" + " task plan has been rolled back to its previous" + " state. Any terminal actions you outputted were not" + " executed. Please review the task goal again and" + " consider what the optimal strategy is now." + ) + ) + + return { + "checker_success": False, + "operator_replied": False, + VALIDATOR_MESSAGES_KEY: [system_msg], + "structured_decisions": "", + } + + else: + logger.info("Planner approved the task plan changes.") + except Exception as e: + logger.error(f"Planner validation task failed: {e}") + + current_step_id = None + if checker_success: + raw_data = state.operator_raw_data + if raw_data: + try: + screenshot_b64 = raw_data.get("screenshot_b64") + xml_hierarchy = raw_data.get("xml_hierarchy") + ocr_results = raw_data.get("ocr_results") + + screenshot_bytes = base64.b64decode(screenshot_b64) + + action_taken = None + if state.structured_decisions: + try: + action_taken = json.loads(state.structured_decisions) + except Exception as e: + logger.warning(f"Failed to parse structured decisions: {e}") + + subgoal_hash, sub_subgoal_hash = _get_active_subgoal_hashes(ctx) + step_id = ctx.data_engine.record_step( + pre_screenshot_bytes=screenshot_bytes, + ui_tree=xml_hierarchy, + ocr_result=ocr_results, + action_taken=action_taken, + operator_raw_thinking=getattr(state, "operator_raw_thinking", None), + operator_native_thinking=getattr(state, "operator_native_thinking", None), + extra_metadata={ + "subgoal_hash": subgoal_hash, + "sub_subgoal_hash": sub_subgoal_hash, + "width": raw_data.get("width"), + "height": raw_data.get("height"), + }, + ) + current_step_id = str(step_id) + logger.info(f"Recorded step in DataEngine: {current_step_id}") + + except Exception as e: + logger.error(f"Failed to record step in DataEngine: {e}") + + # Handle optimistic execution rollback/commit + if ctx.data_engine and getattr(ctx, "task_plan_snapshot", None): + snapshot_dir = ctx.task_plan_snapshot + notes_dir = get_notes_dir(ctx.data_engine.base_dir) + + if checker_success: + logger.info("Checker succeeded. Committing optimistic changes (cleaning up snapshot).") + if snapshot_dir.exists(): + await asyncio.to_thread(shutil.rmtree, snapshot_dir) + logger.info(f"Offloaded snapshot deletion at {snapshot_dir} via asyncio.to_thread") + else: + logger.warning("Checker failed. Rolling back optimistic changes.") + if snapshot_dir.exists(): + + def _restore_and_preserve(): + chat_contents = {} + for p in notes_dir.glob("verification_chat_*.json"): + chat_contents[p.name] = p.read_text(encoding="utf-8") + restore_snapshot(snapshot_dir, notes_dir) + for name, content in chat_contents.items(): + (notes_dir / name).write_text(content, encoding="utf-8") + logger.info(f"Restored verification chat file: {name}") + + await asyncio.to_thread(_restore_and_preserve) + logger.info(f"Restored notes from snapshot {snapshot_dir}") + else: + raise FileNotFoundError( + f"Snapshot directory {snapshot_dir} not found for rollback!" + ) + + update = { + "checker_success": checker_success, + "current_step_id": current_step_id, + "operator_replied": False, + } + if checker_success: + update["subagent_calls"] = [] + return update + + +def execution_check_edge( + state: State, +) -> Literal["execute_decisions", "review_subgoals"]: + if state.checker_success: + if state.structured_decisions: + return "execute_decisions" + return "review_subgoals" + return "review_subgoals" + + +def _get_active_subgoal_hashes(ctx: ArtemisContext) -> tuple[str, str | None]: + """Parses task_plan.md to find the active subgoal and sub-subgoal hashes.""" + + if not ctx.data_engine: + return "default", None + + task_plan_path = get_note_file_path(ctx.data_engine.base_dir, "task_plan") + + if not task_plan_path.exists(): + return "default", None + + try: + content = task_plan_path.read_text(encoding="utf-8") + + return get_active_subgoal_hashes(content) + except Exception as e: + logger.error(f"Failed to parse active subgoal: {e}") + + return "default", None + + +def validate_milestones( + content_before: str, content_after: str, similarity_threshold: float = 0.85 +) -> bool: + """Validates top-level milestones to see if async validation is needed. + + Returns True if structural changes or major wording changes ( str: + cleaned = re.sub(r"\W+", " ", s.lower()).strip() + return re.sub(r"\s+", " ", cleaned) + + for b, a in zip(before_milestones, after_milestones): + if b == a: + continue + + norm_b = normalize(b) + norm_a = normalize(a) + + ratio = difflib.SequenceMatcher(None, norm_b, norm_a).ratio() + if ratio < similarity_threshold: + return True + + return False + + +class _CyFunctionDetectorMeta(type): + def __instancecheck__(self, instance): + name = type(instance).__name__ + return ( + name + in ( + "cyfunction", + "cython_function_or_method", + "builtin_function_or_method", + ) + or "cyfunction" in name.lower() + ) + + +class CyFunctionDetector(metaclass=_CyFunctionDetectorMeta): + pass + + +class NoteArgs(BaseModel): + model_config = {"ignored_types": (CyFunctionDetector,)} + key: str = Field(..., description=SAVE_NOTE_ARG_KEY_DESC) + content: str = Field(..., description=SAVE_NOTE_ARG_CONTENT_DESC) + + +def wrap_note_tool(ctx: ArtemisContext, original_tool): + + async def wrapped_note_tool( + key: str, + content: str, + tool_call_id: Annotated[str, InjectedToolCallId] = None, + state: Annotated[State, InjectedState] = None, + ): + """Wrapped tool to intercept task plan updates.""" + if key == "task_plan" and ctx.data_engine: + task_plan_path = get_note_file_path(ctx.data_engine.base_dir, "task_plan") + + content_before = "" + if task_plan_path.exists(): + content_before = task_plan_path.read_text(encoding="utf-8") + + result = await invoke_tool_with_injection( + original_tool, + {"key": key, "content": content}, + tool_call_id, + state, + record_trace=False, + ) + + content_after = content + if task_plan_path.exists(): + content_after = task_plan_path.read_text(encoding="utf-8") + + lines = content_after.splitlines() + top_level_lines = [line for line in lines if line.startswith("- [")] + if top_level_lines and all(line.startswith("- [x]") for line in top_level_lines): + has_incomplete = any(re.match(r"^\s*-\s*\[[ /]\]", line) for line in lines) + if has_incomplete: + logger.info( + "Rejected: Top-level subgoals are complete but" + " incomplete nested subgoals remain." + ) + if task_plan_path.exists(): + task_plan_path.write_text(content_before, encoding="utf-8") + + return Command( + update={ + VALIDATOR_MESSAGES_KEY: [ + SystemMessage( + content=( + "Your changes to the task plan were not" + " applied. Your changes would cause the" + " task to end completely, but there" + " seem to be goals in the task plan" + " that have not been marked as" + " completed. If you believe these goals" + " are already completed, please update" + " their statuses." + ) + ) + ] + } + ) + + threshold = ( + ctx.execution_setup.planner_validation_threshold + if ctx.execution_setup + and hasattr(ctx.execution_setup, "planner_validation_threshold") + else 0.85 + ) + needs_validation = validate_milestones( + content_before, content_after, similarity_threshold=threshold + ) + disable_planner_validation = ( + ctx.execution_setup and ctx.execution_setup.disable_planner_validation + ) + + if needs_validation and not disable_planner_validation: + logger.info( + "Detected significant top-level task plan modification." + " Triggering async Planner Validation." + ) + + initial_goal = ( + state.initial_goal + if state and hasattr(state, "initial_goal") + else "Unknown Goal" + ) + ctx.task_plan_content_before = content_before + + ctx.planner_task = asyncio.create_task( + run_async_planner_validation( + ctx, + initial_goal, + content_before, + content_after, + getattr(state, "operator_raw_thinking", None), + getattr(state, "operator_native_thinking", None), + ) + ) + + if isinstance(result, Command) and VALIDATOR_MESSAGES_KEY in result.update: + msgs = result.update[VALIDATOR_MESSAGES_KEY] + if msgs and hasattr(msgs[0], "content"): + msgs[0].content = ( + f"{msgs[0].content}\n\nWe have applied your changes" + " to the task plan. A background verification task" + " is currently reviewing these changes. If there" + " are any issues with your modifications, you will" + " be notified shortly." + ) + + # Only match top-level subgoals, intentionally ignoring indented sub-subgoals + disable_checker = ctx.execution_setup and ctx.execution_setup.disable_checker + pattern = re.compile(r"^-\s*\[x\]\s*(.*)$", re.MULTILINE) + + before_matches = set(pattern.findall(content_before)) + after_matches = set(pattern.findall(content_after)) + + new_completions = after_matches - before_matches + + if new_completions and not disable_checker: + logger.info(f"Detected new top-level completions: {new_completions}.") + + all_matches_in_order = pattern.findall(content_after) + latest_completion = None + for match in reversed(all_matches_in_order): + if match in new_completions: + latest_completion = match + break + + if latest_completion: + logger.info(f"Triggering Checker for latest completion: {latest_completion}") + raw_perception = getattr(state, "operator_raw_data", None) + ui_hierarchy = getattr(state, "latest_ui_hierarchy", None) + ctx.checker_task = asyncio.create_task( + run_async_check( + ctx, + latest_completion, + raw_perception_data=raw_perception, + latest_ui_hierarchy=ui_hierarchy, + ) + ) + + return result + else: + return await invoke_tool_with_injection( + original_tool, + {"key": key, "content": content}, + tool_call_id, + state, + record_trace=False, + ) + + return StructuredTool.from_function( + func=lambda *a, **kw: None, + coroutine=wrapped_note_tool, + name=original_tool.name, + description=original_tool.description, + args_schema=NoteArgs, + ) + + +class UpdateNoteArgs(BaseModel): + model_config = {"ignored_types": (CyFunctionDetector,)} + key: str = Field(..., description=UPDATE_NOTE_ARG_KEY_DESC) + target: str = Field(..., description=UPDATE_NOTE_ARG_TARGET_DESC) + replacement: str = Field(..., description=UPDATE_NOTE_ARG_REPLACEMENT_DESC) + + +def wrap_update_note_tool(ctx: ArtemisContext, original_tool): + + async def wrapped_update_note( + key: str, + target: str, + replacement: str, + tool_call_id: Annotated[str, InjectedToolCallId] = None, + state: Annotated[State, InjectedState] = None, + ): + """Wrapped tool to intercept task plan updates via update_note.""" + if key == "task_plan" and ctx.data_engine: + task_plan_path = get_note_file_path(ctx.data_engine.base_dir, "task_plan") + + content_before = "" + if task_plan_path.exists(): + content_before = task_plan_path.read_text(encoding="utf-8") + + result = await invoke_tool_with_injection( + original_tool, + {"key": key, "target": target, "replacement": replacement}, + tool_call_id, + state, + record_trace=False, + ) + + content_after = "" + if task_plan_path.exists(): + content_after = task_plan_path.read_text(encoding="utf-8") + + lines = content_after.splitlines() + top_level_lines = [line for line in lines if line.startswith("- [")] + if top_level_lines and all(line.startswith("- [x]") for line in top_level_lines): + has_incomplete = any(re.match(r"^\s*-\s*\[[ /]\]", line) for line in lines) + if has_incomplete: + logger.info( + "Rejected: Top-level subgoals are complete but" + " incomplete nested subgoals remain." + ) + if task_plan_path.exists(): + task_plan_path.write_text(content_before, encoding="utf-8") + + return Command( + update={ + VALIDATOR_MESSAGES_KEY: [ + SystemMessage( + content=( + "Your changes to the task plan were not" + " applied. Your changes would cause the" + " task to end completely, but there" + " seem to be goals in the task plan" + " that have not been marked as" + " completed. If you believe these goals" + " are already completed, please update" + " their statuses." + ) + ) + ] + } + ) + + threshold = ( + ctx.execution_setup.planner_validation_threshold + if ctx.execution_setup + and hasattr(ctx.execution_setup, "planner_validation_threshold") + else 0.85 + ) + needs_validation = validate_milestones( + content_before, content_after, similarity_threshold=threshold + ) + disable_planner_validation = ( + ctx.execution_setup and ctx.execution_setup.disable_planner_validation + ) + + if needs_validation and not disable_planner_validation: + logger.info( + "Detected significant top-level task plan modification." + " Triggering async Planner Validation." + ) + + initial_goal = ( + state.initial_goal + if state and hasattr(state, "initial_goal") + else "Unknown Goal" + ) + ctx.task_plan_content_before = content_before + + ctx.planner_task = asyncio.create_task( + run_async_planner_validation( + ctx, + initial_goal, + content_before, + content_after, + getattr(state, "operator_raw_thinking", None), + getattr(state, "operator_native_thinking", None), + ) + ) + + if isinstance(result, Command) and VALIDATOR_MESSAGES_KEY in result.update: + msgs = result.update[VALIDATOR_MESSAGES_KEY] + if msgs and hasattr(msgs[0], "content"): + msgs[0].content = ( + f"{msgs[0].content}\n\nWe have applied your changes" + " to the task plan. A background verification task" + " is currently reviewing these changes. If there" + " are any issues with your modifications, you will" + " be notified shortly." + ) + + disable_checker = ctx.execution_setup and ctx.execution_setup.disable_checker + pattern = re.compile(r"^-\s*\[x\]\s*(.*)$", re.MULTILINE) + before_matches = set(pattern.findall(content_before)) + after_matches = set(pattern.findall(content_after)) + new_completions = after_matches - before_matches + + if new_completions and not disable_checker: + logger.info(f"Detected new top-level completions: {new_completions}.") + all_matches_in_order = pattern.findall(content_after) + latest_completion = None + for match in reversed(all_matches_in_order): + if match in new_completions: + latest_completion = match + break + + if latest_completion: + logger.info(f"Triggering Checker for latest completion: {latest_completion}") + raw_perception = getattr(state, "operator_raw_data", None) + ui_hierarchy = getattr(state, "latest_ui_hierarchy", None) + ctx.checker_task = asyncio.create_task( + run_async_check( + ctx, + latest_completion, + raw_perception_data=raw_perception, + latest_ui_hierarchy=ui_hierarchy, + ) + ) + + return result + else: + return await invoke_tool_with_injection( + original_tool, + {"key": key, "target": target, "replacement": replacement}, + tool_call_id, + state, + record_trace=False, + ) + + return StructuredTool.from_function( + func=lambda *a, **kw: None, + coroutine=wrapped_update_note, + name=original_tool.name, + description=original_tool.description, + args_schema=UpdateNoteArgs, + ) + + +async def get_graph(ctx: ArtemisContext) -> CompiledStateGraph: + graph_builder = StateGraph(State) + + def convergence_gate(state: State) -> Literal["continue", "end"]: + logger.info("Starting convergence_gate") + + # Check checker result + if hasattr(state, "checker_success") and not state.checker_success: + logger.info("Checker failed, returning to operator for retry.") + return "continue" + + task_plan_content = "" + if ctx.data_engine: + file_path = get_note_file_path(ctx.data_engine.base_dir, "task_plan") + if file_path.exists(): + try: + task_plan_content = file_path.read_text(encoding="utf-8") + except Exception as e: + logger.error(f"Failed to read task plan: {e}") + return "continue" + else: + logger.warning(f"Task plan file not found at {file_path}") + return "continue" + + lines = task_plan_content.splitlines() + + # Check for all completed + top_level_lines = [line for line in lines if line.startswith("- [")] + + if not top_level_lines: + logger.info("No subgoals found in file, ending") + return "end" + + all_done = all(line.startswith("- [x]") for line in top_level_lines) + if all_done: + logger.info("All subgoals are completed, ending the goal") + return "end" + + return "continue" + + # Get native tools + native_wrappers = [ + save_note_wrapper, + read_note_wrapper, + list_notes_wrapper, + update_note_wrapper, + append_note_wrapper, + ] + + native_tools = get_tools_from_wrappers(ctx, native_wrappers) + + ## Define nodes + graph_builder.add_node("planner", PlannerNode(ctx, tools=native_tools)) + + # Wrap native tools for Operator + operator_native_tools = [] + for t in native_tools: + if t.name == "save_note": + operator_native_tools.append(wrap_note_tool(ctx, t)) + elif t.name == "update_note": + operator_native_tools.append(wrap_update_note_tool(ctx, t)) + elif t.name == "append_note": + operator_native_tools.append(wrap_note_tool(ctx, t)) + else: + operator_native_tools.append(t) + + # Get specialized tools for Operator + operator_specialized_wrappers = [ + ask_diagnoser_wrapper, + run_adb_command_wrapper, + manage_task_wrapper, + ask_explorer_wrapper, + ] + if ctx.execution_setup and ctx.execution_setup.enable_committee: + operator_specialized_wrappers.append(ask_committee_wrapper) + + operator_specialized_tools = get_tools_from_wrappers(ctx, operator_specialized_wrappers) + + # Customize descriptions for run_adb_command and manage_task to match exploratory role + for t in operator_specialized_tools: + if t.name == "run_adb_command": + t.description = ( + "[EXPLORER] Executes a shell command directly on the Android" + " mobile device via ADB shell. Use this tool when normal" + " operating tools cannot achieve the goal, do not rely on it." + " You must not use this tool to bypass goals that could be" + " achieved by action tools.\nCan run synchronously or" + " transition to a background task if execution takes longer" + " than WaitMsBeforeAsync.\nSupports persistent environments" + " (environment variables) on the phone across invocations." + ) + elif t.name == "manage_task": + t.description = ( + "[EXPLORER] Manage background ADB shell tasks launched via" + " run_adb_command to monitor or terminate ongoing background" + " commands." + ) + + operator_tools = operator_native_tools + operator_specialized_tools + + if ctx.execution_setup and ctx.execution_setup.video_recording_tools_enabled: + operator_tools.append(get_video_analyzer_tool(ctx, role="operator")) + + graph_builder.add_node("operator", OperatorNode(ctx, tools=operator_tools)) + graph_builder.add_node("validator", ValidatorNode(ctx)) + + graph_builder.add_node("summarizer", SummarizerNode(ctx)) + graph_builder.add_node("execution_check", functools.partial(execution_check_node, ctx=ctx)) + graph_builder.add_node("perception", functools.partial(perception_node, ctx=ctx)) + graph_builder.add_node(node="convergence", action=convergence_node, defer=True) + + ## Linking nodes + graph_builder.add_edge(START, "planner") + graph_builder.add_edge("planner", "convergence") + graph_builder.add_edge("operator", "execution_check") + graph_builder.add_conditional_edges( + "execution_check", + execution_check_edge, + { + "review_subgoals": "convergence", + "execute_decisions": "validator", + }, + ) + graph_builder.add_edge("validator", "summarizer") + graph_builder.add_edge("summarizer", "convergence") + + graph_builder.add_conditional_edges( + source="convergence", + path=convergence_gate, + path_map={ + "continue": "perception", + "end": END, + }, + ) + graph_builder.add_edge("perception", "operator") + + return graph_builder.compile() diff --git a/artemis/graph/perception.py b/artemis/graph/perception.py new file mode 100644 index 0000000..fb528ff --- /dev/null +++ b/artemis/graph/perception.py @@ -0,0 +1,187 @@ +# 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 asyncio +import base64 +import hashlib +import json +from pathlib import Path + +from artemis.context import ArtemisContext +from artemis.controllers.unified_controller import UnifiedMobileController +from artemis.data_engine.trace import trace +from artemis.graph.state import State +from artemis.utils.logger import get_logger +from artemis.utils.ocr_api import perform_ocr +from artemis.utils.ocr_xml_fusion import ( + _crop_image_remove_status_bar, + _detect_status_bar_height, + _map_coordinates_back, + fuse_ocr_with_xml, +) + +logger = get_logger(__name__) + + +def _should_skip_settling(state: State) -> bool: + """Determine if we should skip screen settling based on the last executed action. + + Purely logical, non-UI-drawing, or wait operations don't require dynamic + settling checks. + """ + if not state or not state.structured_decisions: + # First turn or no prior action, skip settling + return True + + try: + actions = json.loads(state.structured_decisions) + if not actions or not isinstance(actions, list): + return True + + last_action = actions[-1] + action_name = last_action.get("action") + + # Actions that are known to have already completed full wait loops or don't render UI + skip_actions = { + "wait_for_delay", + } + + if action_name in skip_actions: + logger.info(f"Last action was '{action_name}'. Skipping dynamic screen settling.") + return True + + return False + except Exception as e: + logger.warning( + f"Failed to parse last action for settling heuristic: {e}." + " Defaulting to settling check." + ) + return False + + +@trace(type="agent", name="perception") +def _check_injected_instruction_file(base_dir: str) -> str | None: + instruction_path = Path(base_dir) / "injected_instruction.json" + if instruction_path.exists(): + with open(instruction_path, encoding="utf-8") as f: + data = json.load(f) + injected_instruction = data.get("instruction") + instruction_path.unlink() + logger.info(f"Read and deleted injected instruction: '{injected_instruction}'") + return injected_instruction + return None + + +async def perception_node(state: State, ctx: ArtemisContext) -> dict: + """Perception node (fast path): captures screenshot, extracts XML hierarchy, and performs OCR. + + Always executes at the start of an operator cycle. + """ + logger.info("Entering Perception Node...") + + if not ctx.data_engine.current_step_id: + ctx.data_engine.allocate_step_id() + logger.info(f"Pre-allocated step ID in Perception Node: {ctx.data_engine.current_step_id}") + + # Check for injected instruction without blocking event loop + injected_instruction = None + if ctx.data_engine and ctx.data_engine.base_dir: + injected_instruction = await asyncio.to_thread( + _check_injected_instruction_file, str(ctx.data_engine.base_dir) + ) + + controller = UnifiedMobileController(ctx) + + # 1. Capture screenshot with hybrid settling strategy + if _should_skip_settling(state): + logger.info("Skipping settling. Retrieving screen data directly...") + device_data = await controller.get_screen_data() + latest_screenshot_b64 = device_data.base64 + xml_hierarchy = device_data.elements + else: + logger.info("Screen settling required. Delaying 0.4s and capturing screen data...") + await asyncio.sleep(0.4) + device_data = await controller.get_screen_data() + latest_screenshot_b64 = device_data.base64 + xml_hierarchy = device_data.elements + + latest_screenshot_bytes = base64.b64decode(latest_screenshot_b64) + + # Update context device dimensions dynamically + if ctx.device: + ctx.device.device_width = device_data.width + ctx.device.device_height = device_data.height + logger.info(f"Updated context device dimensions: {device_data.width}x{device_data.height}") + + # 3. Perform OCR + screen_height = device_data.height + status_bar_height = _detect_status_bar_height(xml_hierarchy, screen_height) + if status_bar_height > 0: + logger.info( + f"Detected status bar height: {status_bar_height}. Cropping screenshot for OCR." + ) + cropped_b64, _, _ = await asyncio.to_thread( + _crop_image_remove_status_bar, + latest_screenshot_b64, + status_bar_height, + ) + raw_ocr_results = await perform_ocr(cropped_b64) + ocr_results = _map_coordinates_back(raw_ocr_results, status_bar_height) + else: + logger.info("No status bar detected. Performing OCR on full screenshot.") + ocr_results = await perform_ocr(latest_screenshot_b64) + + # 4. Fuse OCR with XML + fused_xml = fuse_ocr_with_xml(xml_hierarchy, ocr_results) + + # 5. Save to Data Engine in background thread + image_name = None + screenshot_path = None + if ctx.data_engine: + image_name = hashlib.sha256(latest_screenshot_bytes).hexdigest() + screenshot_path = str(ctx.data_engine.get_image_path(image_name)) + asyncio.create_task( + asyncio.to_thread( + ctx.data_engine.get_or_create_image, + latest_screenshot_bytes, + ui_tree=xml_hierarchy, + ocr_result=ocr_results, + ) + ) + logger.info(f"Offloaded perception data saving to background with image_name: {image_name}") + + # Cache latest perception data on ArtemisContext for downstream zero-latency reuse (e.g. Checker) + ctx.latest_perception_data = { + "screenshot_b64": latest_screenshot_b64, + "latest_ui_hierarchy": fused_xml, + "xml_hierarchy": xml_hierarchy if xml_hierarchy is not None else [], + "ocr_results": ocr_results, + "width": device_data.width, + "height": device_data.height, + } + + # Return state update + update = { + "latest_screenshot": screenshot_path, + "latest_ui_hierarchy": fused_xml, + "operator_raw_data": { + "screenshot_b64": latest_screenshot_b64, + "xml_hierarchy": xml_hierarchy if xml_hierarchy is not None else [], + "ocr_results": ocr_results, + "width": device_data.width, + "height": device_data.height, + }, + "injected_instruction": injected_instruction, + } + return update diff --git a/artemis/graph/state.py b/artemis/graph/state.py new file mode 100644 index 0000000..c65bcf4 --- /dev/null +++ b/artemis/graph/state.py @@ -0,0 +1,134 @@ +# 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. + +from typing import Annotated, Any +from langchain_core.messages import AnyMessage +from langgraph.graph import add_messages +from artemis.config import AgentNode +from artemis.context import ArtemisContext +from artemis.utils.logger import get_logger +from pydantic import BaseModel, ConfigDict + +logger = get_logger(__name__) + + +def take_last(a, b): + """Reducer function keeping the latest value.""" + return b + + +class State(BaseModel): + model_config = ConfigDict( + arbitrary_types_allowed=True, + extra="allow", + ) + + messages: Annotated[list[AnyMessage], "Sequential messages", add_messages] + + remaining_steps: Annotated[int | None, "Remaining steps before the task is completed"] = None + + # planner related keys + initial_goal: Annotated[str, "Initial goal given by the user"] + injected_instruction: Annotated[ + str | None, + "Injected instruction from user for micro-adjustment", + take_last, + ] = None + + # operator related keys (perception) + latest_ui_hierarchy: Annotated[ + list[dict] | None, "Latest UI hierarchy of the device", take_last + ] + latest_screenshot: Annotated[str | None, "Path to the latest screenshot", take_last] + focused_app_info: Annotated[str | None, "Focused app info", take_last] + device_date: Annotated[str | None, "Date of the device", take_last] + + # operator related keys (decisions) + structured_decisions: Annotated[ + str | None, + "Structured decisions made by the operator, for the validator to follow", + take_last, + ] + operator_raw_thinking: Annotated[ + str | None, + "Raw thinking process of the operator", + take_last, + ] = None + operator_native_thinking: Annotated[ + str | None, + "Native/implicit thinking process of the operator", + take_last, + ] = None + short_term_memory: Annotated[ + str | None, + "Short term memory / scratchpad for the operator", + take_last, + ] = None + complete_subgoals_by_ids: Annotated[ + list[str], + "List of subgoal IDs to complete", + take_last, + ] + current_agent: Annotated[str | None, "Current active agent", take_last] = None + operator_replan_reason: Annotated[str | None, "Reason for operator replan", take_last] = None + subgoal_plan: Annotated[list[Any] | None, "Subgoal plan", take_last] = None + operator_tactical_plan: Annotated[list[Any] | None, "Operator tactical plan", take_last] = None + + # validator related keys + validator_messages: Annotated[list[AnyMessage], "Sequential Validator messages", add_messages] + last_execution_result: Annotated[ + dict | None, "Last execution result from validator", take_last + ] = None + current_step_id: Annotated[str | None, "Current step ID from data engine", take_last] = None + operator_raw_data: Annotated[ + dict | None, + "Raw perception data from operator (screenshot, xml, ocr)", + take_last, + ] = None + checker_success: Annotated[ + bool | None, + "True if checker succeeded or not triggered, False if failed", + take_last, + ] = None + subagent_calls: Annotated[ + list[str], "List of sub-agent calls in current session", take_last + ] = [] + operator_replied: Annotated[bool | None, "True if operator replied to checker", take_last] = ( + None + ) + operator_tool_limit_exceeded: Annotated[ + bool | None, + "True if operator exceeded tool call limit in the previous turn", + take_last, + ] = None + indexed_points: Annotated[ + list[list[int]] | None, + "Active coordinate mappings matching indexed text elements", + take_last, + ] = None + indexed_elements: Annotated[ + list[dict] | None, + "Active interactable elements matching indexed text elements," + " containing text, bounds, class, and center coordinates.", + take_last, + ] = None + + async def asanitize_update( + self, + ctx: ArtemisContext, + update: dict, + agent: AgentNode | None = None, + ): + """Sanitizes the state update to ensure it is valid and apply side effect logic where required.""" + return update diff --git a/artemis/interfaces/cli/__init__.py b/artemis/interfaces/cli/__init__.py new file mode 100644 index 0000000..b62f113 --- /dev/null +++ b/artemis/interfaces/cli/__init__.py @@ -0,0 +1,19 @@ +# 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. + +"""ARTEMIS CLI interface package.""" + +from artemis.interfaces.cli.main import app, cli + +__all__ = ["app", "cli"] diff --git a/artemis/interfaces/cli/commands/__init__.py b/artemis/interfaces/cli/commands/__init__.py new file mode 100644 index 0000000..64e75be --- /dev/null +++ b/artemis/interfaces/cli/commands/__init__.py @@ -0,0 +1,15 @@ +# 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. + +"""CLI command modules.""" diff --git a/artemis/interfaces/cli/commands/batch.py b/artemis/interfaces/cli/commands/batch.py new file mode 100644 index 0000000..3185a48 --- /dev/null +++ b/artemis/interfaces/cli/commands/batch.py @@ -0,0 +1,160 @@ +# 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. + +"""Batch tasks execution command (artemis batch).""" + +import asyncio +import json +from pathlib import Path +from typing import Annotated + +from artemis.config import initialize_llm_config +from artemis.sdk import Agent +from artemis.sdk.builders import Builders +from artemis.sdk.types.task import AgentProfile +from artemis.utils.logger import get_logger +from rich.console import Console +from rich.table import Table +import typer + +logger = get_logger(__name__) + + +async def run_batch_tasks( + tasks: list[str], + profile_name: str = "pro", + delay_seconds: float = 5.0, +) -> None: + """Executes a list of automation tasks sequentially.""" + llm_config = initialize_llm_config() + profile = AgentProfile(name="default", llm_config=llm_config) + config = Builders.AgentConfig.with_default_profile(profile).build() + + agent = Agent(config=config) + await agent.init() + + console = Console() + results = [] + + try: + for idx, goal in enumerate(tasks, start=1): + console.rule(f"[bold cyan]Task {idx}/{len(tasks)}: {goal}[/bold cyan]") + status = "SUCCESS" + error_msg = "" + try: + result = await agent.run_task(goal=goal, profile=profile_name) + logger.info(f"Task {idx} completed: {result}") + except Exception as e: + status = "FAILED" + error_msg = str(e) + logger.error(f"Task {idx} failed: {e}") + + results.append({"task": goal, "status": status, "error": error_msg}) + + if idx < len(tasks) and delay_seconds > 0: + await asyncio.sleep(delay_seconds) + finally: + await agent.clean() + + # Print summary table + table = Table(title="Batch Execution Summary") + table.add_column("#", justify="right", style="cyan") + table.add_column("Task Goal", style="white") + table.add_column("Status", justify="center") + table.add_column("Notes", style="red") + + for i, res in enumerate(results, 1): + status_styled = ( + "[bold green]PASS[/bold green]" + if res["status"] == "SUCCESS" + else "[bold red]FAIL[/bold red]" + ) + table.add_row(str(i), res["task"], status_styled, res["error"]) + + console.print(table) + + +def batch_command( + tasks_file: Annotated[ + Path | None, + typer.Option( + "--file", + "-f", + help="Path to a text or JSON file containing task goals (one per line or JSON array).", + ), + ] = None, + goals: Annotated[ + list[str] | None, + typer.Argument( + help="Task goals specified directly on the command line.", + ), + ] = None, + profile: Annotated[ + str, + typer.Option( + "--profile", + "-p", + help="Execution profile to use for all tasks ('flash' or 'pro').", + ), + ] = "pro", + delay: Annotated[ + float, + typer.Option( + "--delay", + "-d", + help="Delay in seconds between successive tasks.", + ), + ] = 5.0, +) -> None: + """Execute multiple automation tasks in sequence.""" + task_list: list[str] = [] + + if tasks_file: + if not tasks_file.exists(): + typer.secho(f"Error: Tasks file '{tasks_file}' not found.", fg=typer.colors.RED) + raise typer.Exit(1) + + content = tasks_file.read_text(encoding="utf-8").strip() + if tasks_file.suffix.lower() == ".json": + try: + parsed = json.loads(content) + if isinstance(parsed, list): + task_list.extend(str(item) for item in parsed) + else: + typer.secho( + "Error: JSON file must contain a list of goals.", fg=typer.colors.RED + ) + raise typer.Exit(1) + except Exception as e: + typer.secho(f"Error parsing JSON file: {e}", fg=typer.colors.RED) + raise typer.Exit(1) + else: + task_list.extend( + [ + line.strip() + for line in content.splitlines() + if line.strip() and not line.startswith("#") + ] + ) + + if goals: + task_list.extend(goals) + + if not task_list: + typer.secho( + "Error: No tasks provided. Use --file or pass goals as arguments.", fg=typer.colors.RED + ) + raise typer.Exit(1) + + asyncio.run(run_batch_tasks(task_list, profile_name=profile, delay_seconds=delay)) diff --git a/artemis/interfaces/cli/commands/bench.py b/artemis/interfaces/cli/commands/bench.py new file mode 100644 index 0000000..5d8d6d5 --- /dev/null +++ b/artemis/interfaces/cli/commands/bench.py @@ -0,0 +1,157 @@ +# 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. + +"""Benchmark and evaluation command (artemis bench).""" + +import asyncio +from pathlib import Path +import subprocess +import sys +from typing import Annotated + +from artemis.config import initialize_llm_config +from artemis.sdk import Agent +from artemis.sdk.builders import Builders +from artemis.sdk.types.task import AgentProfile +from artemis.utils.logger import get_logger +import typer + +logger = get_logger(__name__) + + +def check_and_start_emulator(adb_path: str = "adb") -> bool: + """Checks if Android emulator is running; optionally starts it.""" + try: + res = subprocess.run([adb_path, "devices"], capture_output=True, text=True, timeout=5) + for line in res.stdout.splitlines(): + if "emulator-" in line and "device" in line: + return True + except Exception: + pass + return False + + +async def run_benchmark( + task_name: str | None = None, + profile_name: str = "pro", + console_port: int = 5554, + adb_path: str = "adb", +) -> None: + """Runs AndroidWorld benchmark task evaluations.""" + # Add android_world-main if present + workspace_root = Path(__file__).resolve().parent.parent.parent.parent.parent + aw_path = workspace_root / "android_world-main" + if aw_path.exists() and str(aw_path) not in sys.path: + sys.path.insert(0, str(aw_path)) + + try: + from android_world import registry + from android_world.env import env_launcher + except ImportError: + logger.error("AndroidWorld benchmark suite is not installed or available in path.") + raise typer.Exit(1) + + task_registry = registry.TaskRegistry() + aw_registry = task_registry.get_registry(task_registry.ANDROID_WORLD_FAMILY) + + if task_name and task_name not in aw_registry: + logger.error(f"Task '{task_name}' not found in AndroidWorld registry.") + available = sorted(list(aw_registry.keys())) + logger.info(f"Available tasks ({len(available)}): {', '.join(available[:10])}...") + raise typer.Exit(1) + + logger.info(f"Initializing AndroidWorld environment on console port {console_port}...") + env = env_launcher.load_and_setup_env( + console_port=console_port, + emulator_setup=False, + adb_path=adb_path, + ) + + llm_config = initialize_llm_config() + profile = AgentProfile(name="default", llm_config=llm_config) + config = Builders.AgentConfig.with_default_profile(profile).build() + + agent = Agent(config=config) + await agent.init() + + tasks_to_run = [task_name] if task_name else list(aw_registry.keys()) + + try: + for t_name in tasks_to_run: + task_type = aw_registry[t_name] + params = task_type.generate_random_params() + task_instance = task_type(params) + + logger.info(f"\n๐ŸŽฏ [Benchmark] Starting: {t_name} | Goal: {task_instance.goal}") + task_instance.initialize_task(env) + + result = await agent.run_task( + goal=task_instance.goal, + name=f"bench_{t_name}", + profile=profile_name, + ) + success = task_instance.is_successful(env) + logger.info(f"๐Ÿ† Task '{t_name}' Result: {result} | Benchmark Verified: {success}") + finally: + await agent.clean() + env.close() + + +def bench_command( + task: Annotated[ + str | None, + typer.Option( + "--task", + "-t", + help="Specific AndroidWorld task name to evaluate (e.g. Clock_1, Calculator_1).", + ), + ] = None, + profile: Annotated[ + str, + typer.Option( + "--profile", + "-p", + help="Agent profile to evaluate ('flash' or 'pro').", + ), + ] = "pro", + adb_path: Annotated[ + str, + typer.Option( + "--adb-path", + help="Path to ADB executable.", + ), + ] = "adb", + console_port: Annotated[ + int, + typer.Option( + "--port", + help="Emulator console port.", + ), + ] = 5554, +) -> None: + """Run AndroidWorld benchmark task evaluations against the agent.""" + if not check_and_start_emulator(adb_path=adb_path): + typer.secho( + f"Warning: No active emulator detected on port {console_port}. Please start an emulator.", + fg=typer.colors.YELLOW, + ) + + asyncio.run( + run_benchmark( + task_name=task, + profile_name=profile, + console_port=console_port, + adb_path=adb_path, + ) + ) diff --git a/artemis/interfaces/cli/commands/doctor.py b/artemis/interfaces/cli/commands/doctor.py new file mode 100644 index 0000000..aa50922 --- /dev/null +++ b/artemis/interfaces/cli/commands/doctor.py @@ -0,0 +1,250 @@ +# 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. + +"""System, environment, and device diagnostics (artemis doctor).""" + +import shutil +import subprocess +import sys + +from rich.console import Console +from rich.panel import Panel +from rich.table import Table + +from artemis.config import ( + parse_llm_config, + settings, +) + + +def doctor_command() -> None: + """Run diagnostics to inspect system dependencies, device connectivity, and configuration.""" + console = Console() + console.print() + console.print( + Panel( + "[bold cyan]โ˜• Artemis System & Environment Doctor[/bold cyan]\n" + "[dim]Checking prerequisites for autonomous mobile testing...[/dim]", + expand=False, + ) + ) + + table = Table(show_header=True, header_style="bold magenta") + table.add_column("Component", style="dim", width=22) + table.add_column("Status", width=12) + table.add_column("Details & Recommendations") + + all_passed = True + + # 1. Python Version + py_ver = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}" + if sys.version_info >= (3, 12): + table.add_row( + "Python Runtime", + "[bold green]โœ” OK[/bold green]", + f"Python {py_ver} (64-bit)", + ) + else: + all_passed = False + table.add_row( + "Python Runtime", + "[bold red]โœ– Warning[/bold red]", + f"Python {py_ver} detected. Python >= 3.12 is strongly recommended.", + ) + + # 2. ADB Installation & Server + adb_path = shutil.which("adb") + if adb_path: + try: + res = subprocess.run( + [adb_path, "version"], + capture_output=True, + text=True, + timeout=5, + ) + adb_info = res.stdout.splitlines()[0] if res.stdout else "Installed" + table.add_row( + "Android ADB", "[bold green]โœ” OK[/bold green]", f"{adb_path} ({adb_info})" + ) + except Exception: + table.add_row("Android ADB", "[bold green]โœ” OK[/bold green]", f"{adb_path}") + else: + all_passed = False + table.add_row( + "Android ADB", + "[bold red]โœ– Missing[/bold red]", + "adb not found in PATH. Install Android Platform-Tools.", + ) + + # 3. Connected Android Devices + devices_info: list[tuple[str, str, str]] = [] + if adb_path: + try: + res = subprocess.run( + [adb_path, "devices", "-l"], + capture_output=True, + text=True, + timeout=5, + ) + lines = [line_str.strip() for line_str in res.stdout.splitlines() if line_str.strip()] + for line in lines[1:]: # skip 'List of devices attached' + parts = line.split() + if len(parts) >= 2: + serial = parts[0] + state = parts[1] + extra = " ".join(parts[2:]) + devices_info.append((serial, state, extra)) + except Exception: + pass + + if devices_info: + active_devices = [d for d in devices_info if d[1] == "device"] + if active_devices: + summary = "\n".join( + [ + f"โ€ข [bold cyan]{d[0]}[/bold cyan] ({d[2] or 'Authorized'})" + for d in active_devices + ] + ) + table.add_row("Android Devices", "[bold green]โœ” Connected[/bold green]", summary) + else: + all_passed = False + table.add_row( + "Android Devices", + "[bold yellow]โš  Unauthorized[/bold yellow]", + "Device detected but unauthorized. Check phone screen for USB Debugging permission prompt.", + ) + else: + table.add_row( + "Android Devices", + "[bold yellow]โš  None Found[/bold yellow]", + "No Android device or emulator detected. Connect phone via USB or start emulator.", + ) + + # 4. LLM & Vision OCR Key Configuration + gemini_key = settings.get_api_key("google") + ocr_key = settings.get_api_key("ocr") + openai_key = settings.get_api_key("openai") + claude_key = settings.get_api_key("anthropic") + openrouter_key = settings.get_api_key("openrouter") + + configured_keys = [] + if gemini_key: + val = gemini_key.get_secret_value() + masked = f"{val[:6]}...{val[-4:]}" if len(val) > 10 else "***" + configured_keys.append(f"Gemini LLM ({masked})") + if openai_key: + configured_keys.append("OpenAI") + if claude_key: + configured_keys.append("Anthropic") + if openrouter_key: + configured_keys.append("OpenRouter") + + if configured_keys: + table.add_row( + "LLM Providers", "[bold green]โœ” Configured[/bold green]", ", ".join(configured_keys) + ) + else: + all_passed = False + table.add_row( + "LLM Providers", + "[bold red]โœ– Missing[/bold red]", + "No LLM key found in .env. Run [bold cyan]artemis init[/bold cyan] to configure in 10 seconds.", + ) + + # 5. Vision OCR Key Check + if ocr_key: + val = ocr_key.get_secret_value() + masked = f"{val[:6]}...{val[-4:]}" if len(val) > 10 else "***" + table.add_row( + "Vision OCR", + "[bold green]โœ” Configured[/bold green]", + f"Google Cloud Vision API ({masked})", + ) + else: + table.add_row( + "Vision OCR", + "[bold yellow]โš  Missing[/bold yellow]", + "OCR_API_KEY / API_KEY not set. Perception fast path will use fallback.", + ) + + # 6. Configuration File Health + try: + cfg = parse_llm_config() + planner_model = cfg.planner.model + table.add_row( + "Configuration", + "[bold green]โœ” Valid[/bold green]", + f"config/artemis.jsonc parsed (Planner: {planner_model})", + ) + except Exception as e: + all_passed = False + table.add_row( + "Configuration", + "[bold red]โœ– Error[/bold red]", + f"Failed to parse config/artemis.jsonc: {e}", + ) + + # 7. Video Recording Tools (FFmpeg & scrcpy) + ffmpeg_path = shutil.which("ffmpeg") + scrcpy_path = shutil.which("scrcpy") + + if ffmpeg_path: + table.add_row( + "FFmpeg Tools", + "[bold green]โœ” Installed[/bold green]", + f"{ffmpeg_path} (Video encoding/trimming ready)", + ) + else: + table.add_row( + "FFmpeg Tools", + "[dim]โšช Optional[/dim]", + "Not found. Install ffmpeg for video compression and audio analysis.", + ) + + if scrcpy_path: + table.add_row( + "Screen Recorder", + "[bold green]โœ” Installed[/bold green]", + f"{scrcpy_path} (Live screen recording ready)", + ) + else: + table.add_row( + "Screen Recorder", + "[dim]โšช Optional[/dim]", + "scrcpy not found. Install scrcpy for --with-video-recording-tools.", + ) + + console.print(table) + console.print() + + if all_passed: + console.print( + Panel( + "[bold green]๐ŸŽ‰ All system checks passed![/bold green]\n\n" + "Run your first automation task now:\n" + ' [bold cyan]artemis run "Open Settings and check Battery level"[/bold cyan]', + title="Status: Ready", + expand=False, + ) + ) + else: + console.print( + Panel( + "[bold yellow]๐Ÿ’ก Tip:[/bold yellow] Run [bold cyan]artemis init[/bold cyan] to quickly set up your API key and device.", + title="Action Required", + expand=False, + ) + ) + console.print() diff --git a/artemis/interfaces/cli/commands/init.py b/artemis/interfaces/cli/commands/init.py new file mode 100644 index 0000000..d22ae3f --- /dev/null +++ b/artemis/interfaces/cli/commands/init.py @@ -0,0 +1,150 @@ +# 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. + +"""Interactive setup wizard for first-time ARTEMIS onboarding (artemis init).""" + +import shutil +import subprocess + +from rich.console import Console +from rich.panel import Panel +from rich.prompt import Prompt + +from artemis.config import ROOT_DIR + + +def init_command() -> None: + """Interactive quickstart wizard to configure API keys and connect your device.""" + console = Console() + console.print() + console.print( + Panel( + "[bold cyan]โ˜• Welcome to Artemis Mobile Testing Framework[/bold cyan]\n" + "[dim]Let's get you set up in 10 seconds![/dim]", + expand=False, + ) + ) + + # 1. Select Provider + console.print("\n[bold]Step 1: Choose Primary LLM Provider[/bold]") + console.print(" [1] Google Gemini (Recommended - Default)") + console.print(" [2] OpenAI (GPT-4o)") + console.print(" [3] Anthropic (Claude 3.5 Sonnet)") + console.print(" [4] OpenRouter") + + choice = Prompt.ask("Select provider", choices=["1", "2", "3", "4"], default="1") + provider_map = { + "1": ("GEMINI_API_KEY", "Google Gemini"), + "2": ("OPENAI_API_KEY", "OpenAI"), + "3": ("ANTHROPIC_API_KEY", "Anthropic"), + "4": ("OPEN_ROUTER_API_KEY", "OpenRouter"), + } + env_key, provider_name = provider_map[choice] + + # 2. Enter API Key + console.print(f"\n[bold]Step 2: Enter your {provider_name} API Key[/bold]") + api_key = Prompt.ask(f"Enter {env_key}", password=True) + while not api_key.strip(): + console.print("[bold red]API Key cannot be empty.[/bold red]") + api_key = Prompt.ask(f"Enter {env_key}", password=True) + + # 3. Scan for Android Devices + console.print("\n[bold]Step 3: Detecting Android Devices / Emulators...[/bold]") + detected_devices = [] + adb_path = shutil.which("adb") + if adb_path: + try: + res = subprocess.run( + [adb_path, "devices", "-l"], + capture_output=True, + text=True, + timeout=5, + ) + lines = [line_str.strip() for line_str in res.stdout.splitlines() if line_str.strip()] + for line in lines[1:]: + parts = line.split() + if len(parts) >= 2 and parts[1] == "device": + serial = parts[0] + extra = " ".join(parts[2:]) + detected_devices.append((serial, extra)) + except Exception: + pass + + selected_serial = "" + if detected_devices: + console.print( + f"[bold green]โœ” Found {len(detected_devices)} connected device(s):[/bold green]" + ) + for idx, (serial, extra) in enumerate(detected_devices, 1): + console.print(f" [{idx}] {serial} ({extra})") + if len(detected_devices) == 1: + selected_serial = detected_devices[0][0] + console.print(f"[dim]Auto-selected device: {selected_serial}[/dim]") + else: + dev_choice = Prompt.ask( + "Select device index", + choices=[str(i) for i in range(1, len(detected_devices) + 1)], + default="1", + ) + selected_serial = detected_devices[int(dev_choice) - 1][0] + else: + console.print( + "[bold yellow]โš  No connected device detected right now.[/bold yellow]\n" + "[dim]Artemis will auto-detect your device when you connect it later.[/dim]" + ) + + # 4. Generate .env File + env_content = [ + "# ==============================================================================", + "# โ˜• ARTEMIS Environment Configuration (Generated by 'artemis init')", + "# ==============================================================================", + f"{env_key}={api_key.strip()}", + "", + "# Android Device Configuration", + "ADB_HOST=127.0.0.1", + "ADB_PORT=5037", + ] + if selected_serial: + env_content.append(f"ADB_DEVICE_SERIAL={selected_serial}") + else: + env_content.append("# ADB_DEVICE_SERIAL=emulator-5554") + + env_content.extend( + [ + "", + "# Execution Defaults", + "ARTEMIS_DEFAULT_PROFILE=pro", + "ARTEMIS_DEFAULT_MODEL=gemini-2.5-flash", + "ARTEMIS_TRACES_DIR=./traces", + "", + ] + ) + + env_path = ROOT_DIR / ".env" + env_path.write_text("\n".join(env_content), encoding="utf-8") + + console.print() + console.print( + Panel( + f"[bold green]๐ŸŽ‰ Setup completed successfully![/bold green]\n\n" + f"Configured [bold cyan]{provider_name}[/bold cyan] in [dim]{env_path}[/dim]\n\n" + "Try running your first task:\n" + ' [bold cyan]artemis run "Open YouTube and search for Lo-Fi Hip Hop"[/bold cyan]\n\n' + "Or run health diagnostics anytime with:\n" + " [bold cyan]artemis doctor[/bold cyan]", + title="โ˜• Ready to Go!", + expand=False, + ) + ) + console.print() diff --git a/artemis/interfaces/cli/commands/mcp.py b/artemis/interfaces/cli/commands/mcp.py new file mode 100644 index 0000000..4cd4560 --- /dev/null +++ b/artemis/interfaces/cli/commands/mcp.py @@ -0,0 +1,42 @@ +# 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. + +"""Model Context Protocol (MCP) server command (artemis mcp).""" + +from typing import Annotated + +from artemis.mcp.adb_server import mcp as adb_mcp +from artemis.utils.logger import get_logger +import typer + +logger = get_logger(__name__) + + +def mcp_command( + server_type: Annotated[ + str, + typer.Option( + "--type", + "-t", + help="Type of MCP server to start ('adb' or 'xml').", + ), + ] = "adb", +) -> None: + """Launch a ARTEMIS Model Context Protocol (MCP) server.""" + if server_type.lower() == "adb": + logger.info("Starting Artemis ADB MCP Server over stdio...") + adb_mcp.run(transport="stdio") + else: + logger.error(f"Unsupported MCP server type: {server_type}") + raise typer.Exit(1) diff --git a/artemis/interfaces/cli/commands/run.py b/artemis/interfaces/cli/commands/run.py new file mode 100644 index 0000000..184d7ce --- /dev/null +++ b/artemis/interfaces/cli/commands/run.py @@ -0,0 +1,335 @@ +# 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. + +"""Task execution command (artemis run).""" + +import asyncio +import os +from pathlib import Path +from shutil import which +from typing import Annotated + +from adbutils import AdbClient +from langchain_core.callbacks.base import Callbacks +from artemis.config import initialize_llm_config, settings +from artemis import Agent, Builders +from artemis.sdk.types.task import AgentProfile +from artemis.utils.cli_helpers import display_device_status +from artemis.utils.logger import get_logger +from artemis.utils.video import check_ffmpeg_available +from rich.console import Console +from rich.panel import Panel +import typer + +logger = get_logger(__name__) + + +async def execute_task( + goal: str, + locked_app_package: str | None = None, + test_name: str | None = None, + traces_output_path_str: str | None = None, + output_description: str | None = None, + graph_config_callbacks: Callbacks = [], + video_recording_tools_enabled: bool | None = None, + profile: str | None = None, + app_path: str | None = None, + enable_planner_validation: bool | None = None, + enable_committee: bool | None = None, + enable_checker: bool | None = None, + enable_step_summarizer: bool | None = None, + enable_outputter: bool | None = None, + force_output_synthesis: bool | None = None, + explorer_version: str | None = None, + explorer_flash_mode: str | None = None, + explorer_pro_mode: str | None = None, +) -> None: + """Executes a single mobile automation task end-to-end. + + Args: + goal: Target objective to achieve on the mobile device. + locked_app_package: Optional package name to constrain actions to. + test_name: Optional test identifier for trace directory naming. + traces_output_path_str: Destination path for recording traces. + output_description: Structured natural language output description. + graph_config_callbacks: LangGraph callback handlers. + video_recording_tools_enabled: Whether screen video analyzer is enabled. + profile: Execution profile ('flash' for fast reactive loop, 'pro' for full graph). + app_path: Optional local APK path to install prior to execution. + enable_planner_validation: Explicit override to enable/disable async planner validation. + enable_committee: Explicit override to enable/disable Multi-Agent Committee tool. + explorer_version: Override default Explorer version mode ('flash', 'pro', or 'ultra'). + explorer_flash_mode: Override Explorer version mode for Flash execution profile. + explorer_pro_mode: Override Explorer version mode for Pro execution profile. + """ + llm_config = initialize_llm_config() + agent_profile = AgentProfile(name="default", llm_config=llm_config) + config = Builders.AgentConfig.with_default_profile(profile=agent_profile) + + if video_recording_tools_enabled is not None: + config.with_video_recording_tools(enabled=video_recording_tools_enabled) + + if enable_planner_validation is not None: + config.with_planner_validation(enabled=enable_planner_validation) + + if enable_committee is not None: + config.with_committee(enabled=enable_committee) + + if enable_checker is not None: + config.with_checker(enabled=enable_checker) + + if enable_step_summarizer is not None: + config.with_flash_step_summarizer(enabled=enable_step_summarizer) + + if enable_outputter is not None or force_output_synthesis is not None: + config.with_outputter( + enabled=enable_outputter if enable_outputter is not None else True, + force_synthesis=bool(force_output_synthesis), + ) + + if ( + explorer_version is not None + or explorer_flash_mode is not None + or explorer_pro_mode is not None + ): + config.with_explorer( + version=explorer_version, + flash_mode=explorer_flash_mode, + pro_mode=explorer_pro_mode, + ) + + if settings.ADB_HOST: + config.with_adb_server(host=settings.ADB_HOST, port=settings.ADB_PORT) + + if graph_config_callbacks: + config.with_graph_config_callbacks(graph_config_callbacks) + + agent: Agent | None = None + try: + agent = Agent(config=config.build()) + await agent.init( + retry_count=int(os.getenv("ARTEMIS_HEALTH_RETRIES", 5)), + retry_wait_seconds=int(os.getenv("ARTEMIS_HEALTH_DELAY", 2)), + ) + + task = agent.new_task(goal) + if locked_app_package: + task.with_locked_app_package(locked_app_package) + if test_name: + trace_path = traces_output_path_str or str(settings.TRACES_PATH) + task.with_name(test_name).with_trace_recording(path=trace_path) + if output_description: + task.with_output_description(output_description) + if profile: + task.using_profile(profile) + if app_path: + task.with_app_path(Path(app_path)) + + llm_result_path = os.getenv("RESULTS_OUTPUT_PATH", None) + if llm_result_path: + task.with_llm_output_saving(path=llm_result_path) + + await agent.run_task(request=task.build()) + finally: + if agent is not None: + await agent.clean() + + +def run_command( + goal: Annotated[str, typer.Argument(help="The main goal for the agent to achieve.")], + profile: Annotated[ + str | None, + typer.Option( + "--profile", + "-p", + help="Execution profile ('flash' for fast reactive, 'pro' for full graph).", + ), + ] = "pro", + locked_app_package: Annotated[ + str | None, + typer.Option( + "--locked-app", + "-a", + help="Optional Android package name to restrict the agent to.", + ), + ] = None, + test_name: Annotated[ + str | None, + typer.Option( + "--test-name", + "-n", + help="Name of the test run for trace recording.", + ), + ] = None, + traces_path: Annotated[ + str | None, + typer.Option( + "--traces-path", + "-t", + help="Directory where execution traces are persisted.", + ), + ] = None, + output_description: Annotated[ + str | None, + typer.Option( + "--output-description", + "-o", + help="Natural language or schema description of the expected output.", + ), + ] = None, + with_video_recording_tools: Annotated[ + bool | None, + typer.Option( + "--with-video-recording-tools/--without-video-recording-tools", + help=( + "Enable or disable dynamic video recording and screen analysis " + "tools (auto-detected if omitted)." + ), + ), + ] = None, + app_path: Annotated[ + str | None, + typer.Option( + "--app-path", + help="Local APK path to install before starting the task.", + ), + ] = None, + enable_planner_validation: Annotated[ + bool | None, + typer.Option( + "--enable-planner-validation/--disable-planner-validation", + help="Enable or disable async Planner validation when task plan milestones change.", + ), + ] = None, + enable_committee: Annotated[ + bool | None, + typer.Option( + "--enable-committee/--disable-committee", + help="Enable or disable Multi-Agent Committee council debate tool for Operator.", + ), + ] = None, + enable_checker: Annotated[ + bool | None, + typer.Option( + "--enable-checker/--disable-checker", + help="Enable or disable visual subgoal verification and rollback by the Checker agent.", + ), + ] = None, + enable_step_summarizer: Annotated[ + bool | None, + typer.Option( + "--enable-step-summarizer/--disable-step-summarizer", + help="Enable or disable Flash asynchronous objective visual step state summarizer.", + ), + ] = None, + enable_outputter: Annotated[ + bool | None, + typer.Option( + "--enable-outputter/--disable-outputter", + help="Enable or disable Outputter post-execution report and structured output synthesis.", + ), + ] = None, + force_output_synthesis: Annotated[ + bool | None, + typer.Option( + "--force-output-synthesis/--no-force-output-synthesis", + help="Force Outputter report synthesis even if no structured output schema is specified.", + ), + ] = None, + explorer_version: Annotated[ + str | None, + typer.Option( + "--explorer-version", + help="Default Explorer version mode ('flash' for 1-shot detection, 'pro' for 3-turn ReAct, 'ultra' for deep pixel reasoning).", + ), + ] = None, + explorer_flash_mode: Annotated[ + str | None, + typer.Option( + "--explorer-flash-mode", + help="Explorer version mode when running under Flash profile (FlashRunner).", + ), + ] = None, + explorer_pro_mode: Annotated[ + str | None, + typer.Option( + "--explorer-pro-mode", + help="Explorer version mode when running under Pro profile (Operator/Validator).", + ), + ] = None, +) -> None: + """Run an autonomous UI automation task on the connected Android device.""" + if with_video_recording_tools: + check_ffmpeg_available() + + console = Console() + + adb_client = None + try: + if which("adb"): + adb_client = AdbClient( + host=settings.ADB_HOST or "localhost", + port=settings.ADB_PORT or 5037, + ) + except Exception: + pass + + display_device_status(console, adb_client=adb_client) + + cancelled = False + try: + asyncio.run( + execute_task( + goal=goal, + locked_app_package=locked_app_package, + test_name=test_name, + traces_output_path_str=traces_path, + output_description=output_description, + video_recording_tools_enabled=with_video_recording_tools, + profile=profile, + app_path=app_path, + enable_planner_validation=enable_planner_validation, + enable_committee=enable_committee, + enable_checker=enable_checker, + enable_step_summarizer=enable_step_summarizer, + enable_outputter=enable_outputter, + force_output_synthesis=force_output_synthesis, + explorer_version=explorer_version, + explorer_flash_mode=explorer_flash_mode, + explorer_pro_mode=explorer_pro_mode, + ) + ) + except KeyboardInterrupt: + cancelled = True + except Exception as e: + err_msg = str(e) + if "API_KEY" in err_msg or "requires" in err_msg: + console.print() + console.print( + Panel( + f"[bold red]โœ– Authentication Error:[/bold red] {err_msg}\n\n" + "๐Ÿ’ก [bold cyan]Quick Fix:[/bold cyan] Run [bold green]artemis init[/bold green] to configure your API key in 10 seconds,\n" + "or add [bold]GEMINI_API_KEY=your_key[/bold] to [dim].env[/dim].", + title="Missing API Key", + expand=False, + ) + ) + console.print() + raise SystemExit(1) + else: + console.print(f"[bold red]Task execution failed:[/bold red] {e}") + raise + finally: + if cancelled: + raise SystemExit(130) diff --git a/artemis/interfaces/cli/commands/server.py b/artemis/interfaces/cli/commands/server.py new file mode 100644 index 0000000..2c6bf40 --- /dev/null +++ b/artemis/interfaces/cli/commands/server.py @@ -0,0 +1,137 @@ +# 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. + +"""Server and cloud proxy commands (artemis server).""" + +import http.server +import socketserver +import subprocess +import time +from typing import Annotated +import urllib.error +import urllib.request + +from artemis.utils.logger import get_logger +import typer + +logger = get_logger(__name__) +server_app = typer.Typer(help="Cloud proxy and server management commands.") + +cached_token = "" +token_expire_time = 0.0 + + +def get_identity_token() -> str: + """Fetch Google Cloud IAM identity token for authenticated requests.""" + global cached_token, token_expire_time + now = time.time() + if not cached_token or now >= token_expire_time: + try: + res = subprocess.check_output( + ["gcloud", "auth", "print-identity-token"], text=True + ).strip() + cached_token = res + token_expire_time = now + 1800 # cache 30 mins + except Exception as e: + logger.warning(f"Failed to get GCP identity token: {e}") + return "" + return cached_token + + +class CloudRunProxyHandler(http.server.BaseHTTPRequestHandler): + target_url: str = "https://artemis-cloud-brain-971929466682.us-central1.run.app" + + def log_message(self, format, *args): + logger.info(f"[{self.log_date_time_string()}] {format % args}") + + def do_GET(self): + self.proxy_request("GET") + + def do_POST(self): + self.proxy_request("POST") + + def do_DELETE(self): + self.proxy_request("DELETE") + + def do_OPTIONS(self): + self.send_response(200) + self.send_header("Access-Control-Allow-Origin", "*") + self.send_header("Access-Control-Allow-Methods", "GET, POST, DELETE, OPTIONS") + self.send_header("Access-Control-Allow-Headers", "*") + self.end_headers() + + def proxy_request(self, method: str): + url = f"{self.target_url}{self.path}" + headers = {} + for k, v in self.headers.items(): + if k.lower() not in ("host", "authorization", "content-length"): + headers[k] = v + + token = get_identity_token() + if token: + headers["Authorization"] = f"Bearer {token}" + headers["X-Artemis-Token"] = "alice-token" + + content_length = int(self.headers.get("Content-Length", 0)) + body = self.rfile.read(content_length) if content_length > 0 else None + + req = urllib.request.Request(url, data=body, headers=headers, method=method) + try: + with urllib.request.urlopen(req, timeout=30) as resp: + self.send_response(resp.status) + for k, v in resp.getheaders(): + if k.lower() not in ("transfer-encoding", "content-encoding", "content-length"): + self.send_header(k, v) + data = resp.read() + self.send_header("Content-Length", str(len(data))) + self.send_header("Access-Control-Allow-Origin", "*") + self.end_headers() + self.wfile.write(data) + except urllib.error.HTTPError as e: + self.send_response(e.code) + for k, v in e.headers.items(): + if k.lower() not in ("transfer-encoding", "content-encoding", "content-length"): + self.send_header(k, v) + data = e.read() + self.send_header("Content-Length", str(len(data))) + self.send_header("Access-Control-Allow-Origin", "*") + self.end_headers() + self.wfile.write(data) + except Exception as e: + self.send_response(500) + self.end_headers() + self.wfile.write(f"Proxy Error: {e}".encode()) + + +@server_app.command("web") +def web_proxy_command( + port: Annotated[int, typer.Option("--port", "-p", help="Local port to bind.")] = 8080, + target: Annotated[ + str, + typer.Option("--target", "-t", help="Target Cloud Run URL."), + ] = "https://artemis-cloud-brain-971929466682.us-central1.run.app", +) -> None: + """Start the local authenticated proxy for ARTEMIS Cloud Web Dashboard.""" + CloudRunProxyHandler.target_url = target + typer.secho( + f"๐Ÿš€ Starting Artemis Cloud Web Proxy at http://localhost:{port}", fg=typer.colors.GREEN + ) + typer.secho(f"๐Ÿ”— Forwarding authenticated requests to: {target}", fg=typer.colors.CYAN) + + socketserver.TCPServer.allow_reuse_address = True + with socketserver.TCPServer(("", port), CloudRunProxyHandler) as httpd: + try: + httpd.serve_forever() + except KeyboardInterrupt: + typer.secho("\n๐Ÿ›‘ Stopping server.", fg=typer.colors.YELLOW) diff --git a/artemis/interfaces/cli/commands/trace.py b/artemis/interfaces/cli/commands/trace.py new file mode 100644 index 0000000..c2c2002 --- /dev/null +++ b/artemis/interfaces/cli/commands/trace.py @@ -0,0 +1,111 @@ +# 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. + +"""Trace inspection and replay commands (artemis trace).""" + +import datetime +import json +from pathlib import Path +from typing import Annotated + +from artemis.config import settings +from artemis.utils.logger import get_logger +from rich.console import Console +from rich.table import Table +import typer + +logger = get_logger(__name__) +trace_app = typer.Typer(help="Inspect, list, and query task execution traces.") + + +@trace_app.command("list") +def list_traces( + traces_path: Annotated[ + Path | None, + typer.Option("--path", "-p", help="Custom traces directory path."), + ] = None, + limit: Annotated[int, typer.Option("--limit", "-l", help="Max sessions to list.")] = 20, +) -> None: + """List recorded execution trace sessions.""" + base_dir = traces_path or settings.TRACES_PATH + if not base_dir.exists(): + typer.secho(f"No traces found at: {base_dir}", fg=typer.colors.YELLOW) + return + + sessions = [d for d in base_dir.iterdir() if d.is_dir() and not d.name.startswith(".")] + sessions.sort(key=lambda x: x.stat().st_mtime, reverse=True) + + console = Console() + table = Table(title=f"Recorded Traces ({len(sessions)} total)") + table.add_column("Session Name", style="cyan") + table.add_column("Last Modified", style="green") + table.add_column("Artifacts", style="white") + + for s in sessions[:limit]: + artifacts = [] + if (s / "recording.mp4").exists() or (s / "recording.mkv").exists(): + artifacts.append("๐ŸŽฌ Video") + if (s / "steps.json").exists(): + artifacts.append("๐Ÿ“‹ Steps") + if (s / "notes").exists(): + artifacts.append("๐Ÿ“ Notes") + mtime = time_str = time_to_str(s.stat().st_mtime) + table.add_row(s.name, mtime, ", ".join(artifacts) or "Screenshots") + + console.print(table) + + +def time_to_str(ts: float) -> str: + return datetime.datetime.fromtimestamp(ts).strftime("%Y-%m-%d %H:%M:%S") + + +@trace_app.command("view") +def view_trace( + session_name: Annotated[str, typer.Argument(help="Name of the trace session directory.")], + traces_path: Annotated[Path | None, typer.Option("--path", "-p")] = None, +) -> None: + """View step-by-step summary of a recorded trace session.""" + base_dir = traces_path or settings.TRACES_PATH + session_dir = base_dir / session_name + + if not session_dir.exists(): + typer.secho(f"Session '{session_name}' not found in {base_dir}", fg=typer.colors.RED) + raise typer.Exit(1) + + steps_file = session_dir / "steps.json" + console = Console() + + if steps_file.exists(): + try: + steps_data = json.loads(steps_file.read_text(encoding="utf-8")) + table = Table(title=f"Trace Steps: {session_name}") + table.add_column("Step", justify="right", style="cyan") + table.add_column("Action", style="green") + table.add_column("Reasoning / Motivation", style="white") + + for step in steps_data: + table.add_row( + str(step.get("step", "-")), + step.get("action", "-"), + step.get("motivation", step.get("thought", "-")), + ) + console.print(table) + return + except Exception as e: + logger.warning(f"Could not parse steps.json: {e}") + + # If steps.json not present, list files + typer.secho(f"Session files in {session_dir}:", fg=typer.colors.CYAN) + for f in session_dir.iterdir(): + typer.echo(f" - {f.name}") diff --git a/artemis/interfaces/cli/main.py b/artemis/interfaces/cli/main.py new file mode 100644 index 0000000..155c9d5 --- /dev/null +++ b/artemis/interfaces/cli/main.py @@ -0,0 +1,99 @@ +# 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. + +"""ARTEMIS Unified Command Line Interface (CLI).""" + +import importlib.metadata +from typing import Annotated + +from artemis.interfaces.cli.commands.batch import batch_command +from artemis.interfaces.cli.commands.bench import bench_command +from artemis.interfaces.cli.commands.doctor import doctor_command +from artemis.interfaces.cli.commands.init import init_command +from artemis.interfaces.cli.commands.mcp import mcp_command +from artemis.interfaces.cli.commands.run import run_command +from artemis.interfaces.cli.commands.server import server_app +from artemis.interfaces.cli.commands.trace import trace_app +from artemis.utils.logger import get_logger +from rich.console import Console +from rich.panel import Panel +import typer + +logger = get_logger(__name__) + +app = typer.Typer( + name="artemis", + help="โ˜• Artemis: Autonomous Multimodal Android Agent & Testing Framework.", + add_completion=False, + pretty_exceptions_enable=False, + no_args_is_help=True, +) + +# Register subcommands +app.command(name="run", help="Execute an autonomous task on a mobile device.")(run_command) +app.command(name="init", help="Interactive quickstart wizard to configure API keys & device.")( + init_command +) +app.command(name="doctor", help="Check system prerequisites, device status, and configuration.")( + doctor_command +) +app.command(name="batch", help="Execute a batch sequence of automation tasks.")(batch_command) +app.command(name="bench", help="Run AndroidWorld benchmark task evaluations.")(bench_command) +app.command(name="mcp", help="Start the Artemis Model Context Protocol (MCP) server.")(mcp_command) +app.add_typer(server_app, name="server", help="Cloud Run proxy and web dashboard server.") +app.add_typer(trace_app, name="trace", help="Inspect and query execution traces.") + + +def version_callback(value: bool): + if value: + try: + ver = importlib.metadata.version("artemis") + except Exception: + ver = "3.6.3" + console = Console() + console.print( + Panel( + f"[bold cyan]Artemis Agent Platform[/bold cyan] v{ver}\n" + "[dim]Autonomous Multimodal Mobile AI Engine[/dim]", + title="โ˜• Artemis", + expand=False, + ) + ) + raise typer.Exit() + + +@app.callback() +def main_callback( + version: Annotated[ + bool, + typer.Option( + "--version", + "-v", + help="Show Artemis version and exit.", + callback=version_callback, + is_eager=True, + ), + ] = False, +): + """ARTEMIS Autonomous Mobile Agent CLI.""" + pass + + +def cli(): + """Main CLI entrypoint.""" + app() + + +if __name__ == "__main__": + cli() diff --git a/artemis/interfaces/sdk/__init__.py b/artemis/interfaces/sdk/__init__.py new file mode 100644 index 0000000..30af13b --- /dev/null +++ b/artemis/interfaces/sdk/__init__.py @@ -0,0 +1,24 @@ +# 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. + +"""Python SDK for ARTEMIS.""" + +from artemis.interfaces.sdk.client import ArtemisClient +from artemis.interfaces.sdk.task import Task, TaskResult + +__all__ = [ + "ArtemisClient", + "Task", + "TaskResult", +] diff --git a/artemis/interfaces/sdk/client.py b/artemis/interfaces/sdk/client.py new file mode 100644 index 0000000..8835e3c --- /dev/null +++ b/artemis/interfaces/sdk/client.py @@ -0,0 +1,147 @@ +# 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. + +from typing import Any, Literal +from collections.abc import AsyncGenerator, Callable +from artemis.engine.pipeline import Pipeline +from artemis.interfaces.sdk.task import StreamEvent, StreamEventType, TaskResult +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +class ArtemisClient: + """The official high-level developer entrypoint for interacting with ARTEMIS.""" + + def __init__( + self, device_id: str = "default-device", default_profile: Literal["flash", "pro"] = "pro" + ): + self.device_id = device_id + self.default_profile = default_profile + + async def run( + self, + goal: str, + profile: Literal["flash", "pro"] | None = None, + max_turns: int = 30, + **kwargs: Any, + ) -> TaskResult: + """Executes a single autonomous task and returns the final TaskResult.""" + target_profile = profile or self.default_profile + logger.info(f"ArtemisClient running task '{goal}' with profile '{target_profile}'...") + + state = await Pipeline.execute( + goal=goal, + profile=target_profile, + device_id=self.device_id, + ) + + return TaskResult( + trace_id=state.trace_id, + status=state.status.value, + turns=state.current_turn, + error=state.error_message, + ) + + async def stream_run( + self, + goal: str, + profile: Literal["flash", "pro"] | None = None, + max_turns: int = 30, + callbacks: list[Callable[[StreamEvent], Any]] | None = None, + **kwargs: Any, + ) -> AsyncGenerator[StreamEvent, None]: + """Executes an autonomous task while yielding real-time StreamEvents asynchronously.""" + target_profile = profile or self.default_profile + logger.info(f"ArtemisClient streaming task '{goal}' with profile '{target_profile}'...") + + # 1. Emit initial starting event + start_event = StreamEvent( + event_type=StreamEventType.STATUS, + step_number=0, + payload={ + "status": "starting", + "goal": goal, + "profile": target_profile, + "device_id": self.device_id, + }, + ) + if callbacks: + for cb in callbacks: + try: + cb(start_event) + except Exception as e: + logger.debug(f"Callback error: {e}") + yield start_event + + # 2. Execute task pipeline + try: + state = await Pipeline.execute( + goal=goal, + profile=target_profile, + device_id=self.device_id, + ) + + # 3. Emit intermediate step events + for step in state.steps: + ev = StreamEvent( + event_type=StreamEventType.STEP_END, + step_number=step.step_number, + payload={ + "thought": step.thought, + "action": step.action_name, + "params": step.action_params, + "result": step.result, + "duration": step.duration_seconds, + }, + ) + if callbacks: + for cb in callbacks: + try: + cb(ev) + except Exception as e: + logger.debug(f"Callback error: {e}") + yield ev + + # 4. Emit completion event + final_ev = StreamEvent( + event_type=StreamEventType.STATUS, + step_number=state.current_turn, + payload={ + "status": state.status.value, + "turns": state.current_turn, + "error": state.error_message, + }, + ) + if callbacks: + for cb in callbacks: + try: + cb(final_ev) + except Exception as e: + logger.debug(f"Callback error: {e}") + yield final_ev + + except Exception as e: + err_ev = StreamEvent( + event_type=StreamEventType.ERROR, + step_number=0, + payload={"error": str(e)}, + ) + if callbacks: + for cb in callbacks: + try: + cb(err_ev) + except Exception: + pass + yield err_ev diff --git a/artemis/interfaces/sdk/task.py b/artemis/interfaces/sdk/task.py new file mode 100644 index 0000000..e8deb99 --- /dev/null +++ b/artemis/interfaces/sdk/task.py @@ -0,0 +1,58 @@ +# 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 time +from typing import Any, Literal +from pydantic import BaseModel, Field + + +class StreamEventType: + THOUGHT = "thought" + ACTION = "action" + OBSERVATION = "observation" + STEP_START = "step_start" + STEP_END = "step_end" + STATUS = "status" + ERROR = "error" + + +class StreamEvent(BaseModel): + """Real-time intermediate progress event emitted during task execution.""" + + event_type: str = Field( + ..., description="Type of event: thought, action, observation, status, etc." + ) + step_number: int = Field(default=0, description="1-indexed execution step") + payload: dict[str, Any] = Field(default_factory=dict, description="Event data dictionary") + timestamp: float = Field(default_factory=time.time, description="Epoch timestamp") + + +class Task(BaseModel): + """Configures a single autonomous mobile automation request.""" + + goal: str = Field(..., description="High-level goal description") + profile: Literal["flash", "pro"] = Field(default="pro", description="Execution profile") + device_id: str = Field(default="default-device", description="Target device serial") + locked_package: str | None = Field(default=None, description="Package lock restriction") + max_turns: int = Field(default=30, description="Turn cutoff limit") + + +class TaskResult(BaseModel): + """Final structured outcome of a task execution.""" + + trace_id: str + status: str + turns: int + output: Any = None + error: str | None = None diff --git a/artemis/llm/__init__.py b/artemis/llm/__init__.py new file mode 100644 index 0000000..626456d --- /dev/null +++ b/artemis/llm/__init__.py @@ -0,0 +1,29 @@ +# 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. + +"""Dynamic LLM Routing and Provider Integration layer.""" + +from artemis.llm.router import ( + ModelEndpoint, + ModelProvider, + ModelRouter, + default_router, +) + +__all__ = [ + "ModelRouter", + "ModelEndpoint", + "ModelProvider", + "default_router", +] diff --git a/artemis/llm/client_factory.py b/artemis/llm/client_factory.py new file mode 100644 index 0000000..be2cf49 --- /dev/null +++ b/artemis/llm/client_factory.py @@ -0,0 +1,25 @@ +# 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. + +"""LLM Provider Client Factory.""" + +from langchain_core.language_models.chat_models import BaseChatModel +from artemis.llm.router import ModelEndpoint, default_router + + +def get_chat_model(role: str = "operator", endpoint: ModelEndpoint | None = None) -> BaseChatModel: + """Instantiates a configured chat model for the requested role or endpoint.""" + if endpoint: + return default_router.instantiate_model(endpoint) + return default_router.create_chat_model(role) diff --git a/artemis/llm/parser.py b/artemis/llm/parser.py new file mode 100644 index 0000000..a852c4d --- /dev/null +++ b/artemis/llm/parser.py @@ -0,0 +1,32 @@ +# 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. + +"""Structured output and thinking tag parser.""" + +import re +from artemis.llm.schemas import ThinkingBlock + + +def extract_thinking_tags(text: str) -> tuple[str, list[ThinkingBlock]]: + """Extracts content inside or tags from LLM response.""" + thinking_blocks: list[ThinkingBlock] = [] + + pattern = r"<(?:thought|thinking)>(.*?)" + matches = re.findall(pattern, text, flags=re.DOTALL) + + for m in matches: + thinking_blocks.append(ThinkingBlock(content=m.strip(), is_native=False)) + + cleaned_text = re.sub(pattern, "", text, flags=re.DOTALL).strip() + return cleaned_text, thinking_blocks diff --git a/artemis/llm/router.py b/artemis/llm/router.py new file mode 100644 index 0000000..55e8d92 --- /dev/null +++ b/artemis/llm/router.py @@ -0,0 +1,433 @@ +# 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. + +"""Dynamic Multi-Provider Model Router for ARTEMIS. + +Provides role-based LLM/VLM dispatching, dynamic fallback chains, +and unified provider configuration across Gemini, Vertex AI, OpenAI, +Anthropic, OpenRouter, and local Ollama/vLLM endpoints. +""" + +from enum import StrEnum +import os +from typing import Any + +from langchain_core.language_models.chat_models import BaseChatModel +from pydantic import BaseModel, Field + +from artemis.config.settings import settings +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +class ModelProvider(StrEnum): + GOOGLE = "google" + GEMINI = "google" + VERTEX_AI = "vertexai" + OPENAI = "openai" + ANTHROPIC = "anthropic" + OPENROUTER = "openrouter" + XAI = "xai" + OLLAMA = "ollama" + VLLM = "vllm" + CUSTOM = "custom" + + @classmethod + def from_string(cls, val: Any) -> "ModelProvider": + """Converts any string, enum, or mock representation into a canonical ModelProvider.""" + if isinstance(val, cls): + return val + s = str(val).lower().replace("_", "").replace("-", "").strip() + mapping = { + "google": cls.GOOGLE, + "gemini": cls.GOOGLE, + "vertexai": cls.VERTEX_AI, + "vertex": cls.VERTEX_AI, + "openai": cls.OPENAI, + "anthropic": cls.ANTHROPIC, + "claude": cls.ANTHROPIC, + "openrouter": cls.OPENROUTER, + "xai": cls.XAI, + "grok": cls.XAI, + "ollama": cls.OLLAMA, + "vllm": cls.VLLM, + "custom": cls.CUSTOM, + } + return mapping.get(s, cls.GOOGLE) + + +class ModelEndpoint(BaseModel): + """Configuration definition for an LLM/VLM model endpoint.""" + + provider: ModelProvider = Field(default=ModelProvider.GOOGLE, description="Model provider") + model_name: str = Field(default="gemini-2.5-flash", description="Model name identifier") + api_key: str | None = Field(default=None, description="API Key or secret") + api_base: str | None = Field(default=None, description="Custom API endpoint base URL") + temperature: float = Field(default=0.0, description="Sampling temperature") + max_tokens: int | None = Field(default=None, description="Maximum completion tokens") + timeout_seconds: float = Field(default=60.0, description="Request timeout in seconds") + is_multimodal: bool = Field( + default=True, description="Whether endpoint accepts image/video inputs" + ) + reasoning_effort: str | None = Field( + default=None, description="Reasoning effort: 'low', 'medium', 'high'" + ) + thinking_budget: int | None = Field(default=None, description="Thinking budget token count") + thinking_level: str | None = Field( + default=None, description="Thinking level: 'minimal', 'low', 'medium', 'high'" + ) + include_thoughts: bool | None = Field( + default=None, description="Whether to include thought traces in response" + ) + enable_grounding: bool = Field( + default=False, + description="Whether to enable Google Search grounding for Gemini endpoints", + ) + + def cache_key(self) -> tuple: + """Returns a deterministic hashable cache key for connection pooling.""" + return ( + self.provider.value, + self.model_name, + self.temperature, + self.timeout_seconds, + self.thinking_budget, + self.thinking_level, + self.include_thoughts, + self.reasoning_effort, + self.enable_grounding, + self.api_base, + ) + + +def _patch_langchain_google_genai(): + """Patches ChatGoogleGenerativeAI._process_tool_config to preserve and ensure include_server_side_tool_invocations.""" + try: + from langchain_google_genai.chat_models import ChatGoogleGenerativeAI + from google.genai.types import ToolConfig + + if getattr(ChatGoogleGenerativeAI, "_is_latte_patched", False): + return + + original_process = ChatGoogleGenerativeAI._process_tool_config + + def patched_process_tool_config(self, tool_choice, tool_config, formatted_tools): + config = original_process(self, tool_choice, tool_config, formatted_tools) + + has_builtin = False + has_functions = False + if formatted_tools: + for tool in formatted_tools: + tool_dict = tool if isinstance(tool, dict) else getattr(tool, "__dict__", {}) + if any( + tool_dict.get(k) is not None + for k in ( + "google_search", + "code_execution", + "google_maps", + "google_search_retrieval", + ) + ): + has_builtin = True + if tool_dict.get("function_declarations"): + has_functions = True + + if tool_config: + normalized = ( + ToolConfig.model_validate(tool_config) + if isinstance(tool_config, dict) + else tool_config + ) + if getattr(normalized, "include_server_side_tool_invocations", None) is not None: + if config is None: + config = ToolConfig( + include_server_side_tool_invocations=normalized.include_server_side_tool_invocations + ) + else: + config.include_server_side_tool_invocations = ( + normalized.include_server_side_tool_invocations + ) + + if has_builtin and has_functions: + if config is None: + config = ToolConfig(include_server_side_tool_invocations=True) + else: + config.include_server_side_tool_invocations = True + + return config + + ChatGoogleGenerativeAI._process_tool_config = patched_process_tool_config + ChatGoogleGenerativeAI._is_latte_patched = True + except Exception as e: + logger.warning(f"Could not patch ChatGoogleGenerativeAI._process_tool_config: {e}") + + +class ModelFactory: + """Unified factory for instantiating and caching LangChain chat models across providers.""" + + _cache: dict[tuple, BaseChatModel] = {} + + @classmethod + def get_model(cls, endpoint: ModelEndpoint) -> BaseChatModel: + key = endpoint.cache_key() + if key not in cls._cache: + cls._cache[key] = cls.create_model(endpoint) + return cls._cache[key] + + @classmethod + def create_model(cls, endpoint: ModelEndpoint) -> BaseChatModel: + provider = ModelProvider.from_string(endpoint.provider) + + if provider == ModelProvider.GOOGLE: + _patch_langchain_google_genai() + from langchain_google_genai import ( + ChatGoogleGenerativeAI, + HarmBlockThreshold, + HarmCategory, + ) + + api_key = ( + endpoint.api_key + or (settings.GOOGLE_API_KEY.get_secret_value() if settings.GOOGLE_API_KEY else None) + or os.environ.get("GOOGLE_API_KEY") + or os.environ.get("GEMINI_API_KEY") + ) + kwargs: dict[str, Any] = { + "model": endpoint.model_name, + "temperature": endpoint.temperature, + "max_output_tokens": endpoint.max_tokens, + "api_key": api_key, + "timeout": endpoint.timeout_seconds, + "thinking_budget": endpoint.thinking_budget, + "thinking_level": endpoint.thinking_level, + "include_thoughts": endpoint.include_thoughts, + "safety_settings": { + HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE, + HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE, + HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE, + HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE, + }, + } + return ChatGoogleGenerativeAI(**{k: v for k, v in kwargs.items() if v is not None}) + + elif provider == ModelProvider.VERTEX_AI: + from langchain_google_vertexai import ( + ChatVertexAI, + HarmBlockThreshold, + HarmCategory, + ) + + kwargs = { + "model_name": endpoint.model_name, + "temperature": endpoint.temperature, + "max_output_tokens": endpoint.max_tokens, + "timeout": endpoint.timeout_seconds, + "thinking_budget": endpoint.thinking_budget, + "safety_settings": { + HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE, + HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE, + HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE, + HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE, + }, + } + return ChatVertexAI(**{k: v for k, v in kwargs.items() if v is not None}) + + elif provider == ModelProvider.OPENAI: + from langchain_openai import ChatOpenAI + + api_key = ( + endpoint.api_key + or (settings.OPENAI_API_KEY.get_secret_value() if settings.OPENAI_API_KEY else None) + or os.environ.get("OPENAI_API_KEY", "EMPTY") + ) + base_url = endpoint.api_base or ( + str(settings.OPENAI_BASE_URL) if settings.OPENAI_BASE_URL else None + ) + kwargs = { + "model": endpoint.model_name, + "temperature": endpoint.temperature, + "max_tokens": endpoint.max_tokens, + "api_key": api_key, + "base_url": base_url, + "timeout": endpoint.timeout_seconds, + } + if endpoint.reasoning_effort: + kwargs["reasoning_effort"] = endpoint.reasoning_effort + return ChatOpenAI(**{k: v for k, v in kwargs.items() if v is not None}) + + elif provider == ModelProvider.ANTHROPIC: + from langchain_anthropic import ChatAnthropic + + api_key = ( + endpoint.api_key + or ( + settings.ANTHROPIC_API_KEY.get_secret_value() + if settings.ANTHROPIC_API_KEY + else None + ) + or os.environ.get("ANTHROPIC_API_KEY") + ) + kwargs = { + "model": endpoint.model_name, + "temperature": endpoint.temperature, + "api_key": api_key, + "timeout": endpoint.timeout_seconds, + } + budget = endpoint.thinking_budget + if not budget and endpoint.reasoning_effort: + effort_map = {"low": 2048, "medium": 8192, "high": 32768} + budget = effort_map.get(endpoint.reasoning_effort.lower()) + if budget: + kwargs["thinking"] = {"type": "enabled", "budget_tokens": budget} + kwargs["temperature"] = 1.0 + return ChatAnthropic(**{k: v for k, v in kwargs.items() if v is not None}) + + elif provider == ModelProvider.OPENROUTER: + from langchain_openai import ChatOpenAI + + api_key = ( + endpoint.api_key + or ( + settings.OPEN_ROUTER_API_KEY.get_secret_value() + if settings.OPEN_ROUTER_API_KEY + else None + ) + or os.environ.get("OPEN_ROUTER_API_KEY") + ) + return ChatOpenAI( + model=endpoint.model_name, + temperature=endpoint.temperature, + api_key=api_key, + base_url=endpoint.api_base or "https://openrouter.ai/api/v1", + timeout=endpoint.timeout_seconds, + ) + + elif provider == ModelProvider.XAI: + from langchain_openai import ChatOpenAI + + api_key = ( + endpoint.api_key + or (settings.XAI_API_KEY.get_secret_value() if settings.XAI_API_KEY else None) + or os.environ.get("XAI_API_KEY") + ) + return ChatOpenAI( + model=endpoint.model_name, + temperature=endpoint.temperature, + api_key=api_key, + base_url=endpoint.api_base or "https://api.x.ai/v1", + timeout=endpoint.timeout_seconds, + ) + + elif provider in (ModelProvider.OLLAMA, ModelProvider.VLLM, ModelProvider.CUSTOM): + from langchain_openai import ChatOpenAI + + api_key = endpoint.api_key or os.environ.get("OPENAI_API_KEY", "EMPTY") + base_url = endpoint.api_base or os.environ.get( + "OPENAI_BASE_URL", "http://localhost:8000/v1" + ) + kwargs = { + "model": endpoint.model_name, + "temperature": endpoint.temperature, + "max_tokens": endpoint.max_tokens, + "api_key": api_key, + "base_url": base_url, + "timeout": endpoint.timeout_seconds, + } + if endpoint.reasoning_effort: + kwargs["reasoning_effort"] = endpoint.reasoning_effort + return ChatOpenAI(**{k: v for k, v in kwargs.items() if v is not None}) + + else: + raise ValueError(f"Unsupported model provider: {provider}") + + +class ModelRouter: + """Dynamic router mapping agent roles to primary and fallback model endpoints.""" + + def __init__(self, default_model: str = "gemini-2.5-flash"): + self.default_endpoint = ModelEndpoint( + provider=ModelProvider.GOOGLE, + model_name=default_model, + ) + self._role_routes: dict[str, ModelEndpoint] = {} + self._fallback_routes: dict[str, list[ModelEndpoint]] = {} + self._setup_default_routes() + + def _setup_default_routes(self) -> None: + """Initializes recommended role mappings.""" + self._role_routes["planner"] = ModelEndpoint( + provider=ModelProvider.GOOGLE, + model_name="gemini-2.5-pro", + temperature=0.2, + ) + self._role_routes["operator"] = ModelEndpoint( + provider=ModelProvider.GOOGLE, + model_name="gemini-2.5-flash", + temperature=0.0, + ) + self._role_routes["validator"] = ModelEndpoint( + provider=ModelProvider.GOOGLE, + model_name="gemini-2.5-pro", + temperature=0.1, + ) + self._role_routes["summarizer"] = ModelEndpoint( + provider=ModelProvider.GOOGLE, + model_name="gemini-2.5-flash", + temperature=0.0, + ) + self._fallback_routes["planner"] = [ + ModelEndpoint(provider=ModelProvider.GOOGLE, model_name="gemini-2.5-flash"), + ] + + def set_role_route( + self, + role: str, + endpoint: ModelEndpoint | str, + fallbacks: list[ModelEndpoint] | None = None, + ) -> None: + """Configures or overrides a route for a specific agent role.""" + if isinstance(endpoint, str): + endpoint = ModelEndpoint(model_name=endpoint) + self._role_routes[role.lower()] = endpoint + if fallbacks is not None: + self._fallback_routes[role.lower()] = fallbacks + + def get_endpoint(self, role: str) -> ModelEndpoint: + """Retrieves the configured primary model endpoint for a role.""" + return self._role_routes.get(role.lower(), self.default_endpoint) + + def get_fallbacks(self, role: str) -> list[ModelEndpoint]: + """Retrieves fallback endpoints for a role.""" + return self._fallback_routes.get(role.lower(), []) + + def instantiate_model(self, endpoint: ModelEndpoint) -> BaseChatModel: + """Instantiates a concrete LangChain chat model via the unified ModelFactory.""" + return ModelFactory.create_model(endpoint) + + def create_chat_model(self, role: str) -> BaseChatModel: + """Creates a chat model with configured fallbacks for the specified role.""" + primary_endpoint = self.get_endpoint(role) + primary_model = self.instantiate_model(primary_endpoint) + + fallbacks = self.get_fallbacks(role) + if fallbacks: + fallback_models = [self.instantiate_model(ep) for ep in fallbacks] + return primary_model.with_fallbacks(fallback_models) + + return primary_model + + +# Global singleton router instance +default_router = ModelRouter() diff --git a/artemis/llm/schemas.py b/artemis/llm/schemas.py new file mode 100644 index 0000000..51241a7 --- /dev/null +++ b/artemis/llm/schemas.py @@ -0,0 +1,36 @@ +# 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. + +"""Dynamic LLM Schema definitions.""" + +from pydantic import BaseModel, Field + + +class ModelParameters(BaseModel): + """Fine-grained generation hyperparameters.""" + + temperature: float = Field(default=0.0, ge=0.0, le=2.0) + top_p: float | None = Field(default=None, ge=0.0, le=1.0) + top_k: int | None = Field(default=None, ge=1) + max_output_tokens: int = Field(default=4096, ge=1) + seed: int | None = Field(default=None) + + +class ThinkingBlock(BaseModel): + """Extracted internal chain-of-thought from thinking models.""" + + content: str = Field(..., description="Raw thought reasoning string") + is_native: bool = Field( + default=False, description="Whether thought was returned via native API field" + ) diff --git a/artemis/main.py b/artemis/main.py new file mode 100644 index 0000000..2de9eaa --- /dev/null +++ b/artemis/main.py @@ -0,0 +1,41 @@ +# 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. + +"""Backward compatibility entrypoint for artemis CLI.""" + +import sys +from artemis.interfaces.cli.main import app + + +def cli(): + # If invoked directly as python -m artemis.main without subcommand 'run', + # check if first argument is a goal rather than a subcommand + if len(sys.argv) > 1 and sys.argv[1] not in ( + "run", + "init", + "doctor", + "batch", + "bench", + "mcp", + "server", + "trace", + "--help", + "-h", + ): + sys.argv.insert(1, "run") + app() + + +if __name__ == "__main__": + cli() diff --git a/artemis/mcp/adb_server.py b/artemis/mcp/adb_server.py new file mode 100644 index 0000000..0105d88 --- /dev/null +++ b/artemis/mcp/adb_server.py @@ -0,0 +1,465 @@ +# 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 asyncio +import json +import logging +import os +from pathlib import Path +import re +import sys + +# Ensure repository root is in sys.path when executed directly or via MCP runner +_REPO_ROOT = str(Path(__file__).resolve().parent.parent.parent) +if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) + +# pylint: disable=wrong-import-position +from adbutils import AdbClient +from mcp.server.fastmcp import Context, FastMCP + +from artemis.clients.ui_automator_client import UIAutomatorClient +from artemis.context import ArtemisContext, DeviceContext, DevicePlatform +from artemis.controllers.unified_controller import UnifiedMobileController +from artemis.utils.app_launch_utils import launch_app_with_retries + +# Redirect stdio to prevent logs from breaking MCP JSON-RPC protocol and causing deadlocks. +try: + from artemis.config import settings + + log_path = Path(settings.TRACES_PATH) / "mcp_server.log" +except Exception: + log_path = Path("traces") / "mcp_server.log" + +log_path.parent.mkdir(parents=True, exist_ok=True) + +for name in ["artemis", __name__]: + log_instance = logging.getLogger(name) + log_instance.setLevel(logging.DEBUG) + log_instance.propagate = False + # Remove standard streams handlers + log_instance.handlers = [ + h for h in log_instance.handlers if not isinstance(h, logging.StreamHandler) + ] + fh = logging.FileHandler(log_path, encoding="utf-8") + fh.setFormatter(logging.Formatter("%(asctime)s | %(name)s | %(levelname)s | %(message)s")) + log_instance.addHandler(fh) + +logger = logging.getLogger(__name__) +logger.info("MCP Server Python process fully started. Configuring FastMCP service...") + +# Avoid concurrent SQLite WAL locks by blocking child processes from connecting to primary DataEngine. +os.environ["ARTEMIS_IPC_PORT"] = "" + +# Create minimal MCP server +mcp = FastMCP("Android_ADB_Controller") + +_GLOBAL_CONTROLLER = None + + +def _get_controller(): + """Lazy-load device controller on-demand as a singleton.""" + global _GLOBAL_CONTROLLER + if _GLOBAL_CONTROLLER is not None: + return _GLOBAL_CONTROLLER + + logger.info("Initializing lazy device controller...") + if os.environ.get("ARTEMIS_CLOUD_MODE") == "1": + logger.info( + "GCP CLOUD MODE active: initializing CloudMobileDeviceController via UnifiedMobileController..." + ) + target_serial = ( + os.environ.get("ARTEMIS_DEVICE_ID") + or os.environ.get("ADB_DEVICE_SERIAL") + or "cloud_device" + ) + ctx = ArtemisContext( + device=DeviceContext( + platform=DevicePlatform.ANDROID, + device_id=target_serial, + width=1080, + height=2400, + ) + ) + _GLOBAL_CONTROLLER = UnifiedMobileController(ctx=ctx) + return _GLOBAL_CONTROLLER + + host = os.environ.get("ADB_HOST", "localhost") + port_str = os.environ.get("ADB_PORT", "5037") + port = int(port_str) if port_str.isdigit() else 5037 + if "ADB_SERVER_SOCKET" not in os.environ and (host != "localhost" or port != 5037): + os.environ["ADB_SERVER_SOCKET"] = f"tcp:{host}:{port}" + adb = AdbClient(host=host, port=port) + + target_serial = os.environ.get("ARTEMIS_DEVICE_ID") or os.environ.get("ADB_DEVICE_SERIAL") + devices = adb.device_list() + if not devices: + raise Exception(f"No Android devices found at {host}:{port}") + + if target_serial: + matched = [d for d in devices if d.serial == target_serial] + if not matched: + raise Exception( + f"Target Android device '{target_serial}' not found at" + f" {host}:{port} among available devices." + ) + device = matched[0] + else: + device = devices[0] + device_id = device.serial + + ui_client = UIAutomatorClient(device_id=device_id) + try: + ui_data = ui_client.get_screen_data() + width, height = ui_data.width, ui_data.height + except Exception as e: + logger.warning(f"Failed initial screen data check, using defaults: {e}") + width, height = 1080, 2400 + + ctx = ArtemisContext( + trace_id="mcp-session", + device=DeviceContext( + host_platform="LINUX", + mobile_platform=DevicePlatform.ANDROID, + device_id=device_id, + device_width=width, + device_height=height, + ), + adb_client=adb, + ui_adb_client=ui_client, + ) + + _GLOBAL_CONTROLLER = UnifiedMobileController(ctx) + logger.info(f"Lazy device controller fully initialized for device: {device_id}") + return _GLOBAL_CONTROLLER + + +def _find_element_at_coords(elements: list[dict], x: int, y: int) -> dict | None: + """Finds the smallest (leaf-most) focusable element containing [x, y].""" + matching_element = None + min_area = float("inf") + + for elem in elements: + is_focusable = ( + elem.get("focusable") == "true" + or elem.get("clickable") == "true" + or "EditText" in str(elem.get("class", "")) + ) + if not is_focusable: + continue + + bounds_str = elem.get("bounds") + if bounds_str and isinstance(bounds_str, str): + match = re.match(r"\[(\-?\d+),(\-?\d+)\]\[(\-?\d+),(\-?\d+)\]", bounds_str) + if match: + x1, y1, x2, y2 = map(int, match.groups()) + if x1 <= x <= x2 and y1 <= y <= y2: + area = (x2 - x1) * (y2 - y1) + if area < min_area: + min_area = area + matching_element = elem + + return matching_element + + +async def _ensure_focus_at_coords(controller, x: int, y: int) -> str | None: + """Ensures the element at [x, y] is focused, tapping only if it is not already focused.""" + try: + elements = await controller.get_ui_elements() + elem = _find_element_at_coords(elements, x, y) + if elem and elem.get("focused") == "true": + logger.info(f"Element under [{x}, {y}] is already focused. Skipping tap.") + return None + except Exception as e: + logger.warning(f"Failed to check focus status: {e}. Falling back to unconditional tap.") + + result = await controller.tap_at(x=x, y=y) + if hasattr(result, "error") and result.error: + return result.error + + # Wait for the UI to settle and keyboard to pop up after tapping + await asyncio.sleep(1.0) + return None + + +@mcp.tool() +async def tap( + ctx: Context, + coordinates: list[int], + times: int = 1, + delay_ms: int = 100, +) -> str: + """Taps on the screen at coordinates. + + 'coordinates' is a list [x, y]. + 'times' is the number of consecutive clicks (default 1). + 'delay_ms' is the delay in milliseconds between consecutive clicks (default + 100). + """ + try: + controller = _get_controller() + except Exception as e: + return f"Error: Controller lazy initialization failed: {e}" + + if len(coordinates) != 2: + return "Error: coordinates must be [x, y]" + + result = await controller.tap_at( + x=coordinates[0], y=coordinates[1], times=times, delay_ms=delay_ms + ) + + if result.error: + return f"Error: {result.error}" + return "Success" + + +@mcp.tool() +async def long_press_on( + ctx: Context, + coordinates: list[int], + duration: int = 1000, +) -> str: + """Long presses on the screen at coordinates. + + 'coordinates' is a list [x, y]. + 'duration' is in milliseconds (default 1000). + """ + try: + controller = _get_controller() + except Exception as e: + return f"Error: Controller lazy initialization failed: {e}" + + if len(coordinates) != 2: + return "Error: coordinates must be [x, y]" + + result = await controller.tap_at( + x=coordinates[0], + y=coordinates[1], + long_press=True, + long_press_duration=duration, + ) + + if result.error: + return f"Error: {result.error}" + return "Success" + + +@mcp.tool() +async def swipe( + ctx: Context, + coordinates: list[int], + duration: int = 400, +) -> str: + """Swipes from start coordinates to end coordinates. + + 'coordinates' is a list [start_x, start_y, end_x, end_y]. + 'duration' is in milliseconds (default 400). + + Set duration >= 1000 to drag-and-drop. Drag slightly past the target + position to trigger reordering. + """ + try: + controller = _get_controller() + except Exception as e: + return f"Error: Controller lazy initialization failed: {e}" + + if len(coordinates) != 4: + return "Error: coordinates must be [start_x, start_y, end_x, end_y]" + + error = await controller.swipe_coords( + start_x=coordinates[0], + start_y=coordinates[1], + end_x=coordinates[2], + end_y=coordinates[3], + duration=duration, + ) + if error: + return f"Error: {error}" + return "Success" + + +@mcp.tool() +async def back(ctx: Context) -> str: + """Simulates pressing the system back button.""" + try: + controller = _get_controller() + except Exception as e: + return f"Error: Controller lazy initialization failed: {e}" + + success = await controller.go_back() + return "Success" if success else "Failed" + + +@mcp.tool() +async def launch_app(ctx: Context, package_name: str) -> str: + """Launches an application by its Android package name with retries and smart polling.""" + try: + controller = _get_controller() + except Exception as e: + return f"Error: Controller lazy initialization failed: {e}" + + success, error_msg = await launch_app_with_retries(controller.ctx, package_name) + return "Success" if success else f"Failed: {error_msg}" + + +@mcp.tool() +async def stop_app(ctx: Context, package_name: str) -> str: + """Force stops an application by its Android package name.""" + try: + controller = _get_controller() + except Exception as e: + return f"Error: Controller lazy initialization failed: {e}" + + success = await controller.terminate_app(package_name) + return "Success" if success else "Failed" + + +@mcp.tool() +async def open_link(ctx: Context, url: str) -> str: + """Opens a URL or deep link on the device.""" + try: + controller = _get_controller() + except Exception as e: + return f"Error: Controller lazy initialization failed: {e}" + + success = await controller.open_url(url) + return "Success" if success else "Failed" + + +@mcp.tool() +async def focus_and_input_text( + ctx: Context, + coordinates: list[int], + text: str, + clear_before_input: bool = False, +) -> str: + """Focuses on a UI element at coordinates and inputs text. + + 'coordinates' is a list [x, y] to tap first to gain focus. + 'clear_before_input' if True, clears all existing text before typing. If False, appends text at the end of existing content. + 'text' supports multi-line content with '\\n'. + """ + try: + controller = _get_controller() + except Exception as e: + return f"Error: Controller lazy initialization failed: {e}" + + if len(coordinates) != 2: + return "Error: coordinates must be [x, y]" + + err = await _ensure_focus_at_coords(controller, coordinates[0], coordinates[1]) + if err: + return f"Error focusing element: {err}" + + if clear_before_input: + success = await controller.erase_text() + if not success: + return "Failed to clear existing text" + else: + # Move cursor to the end for reliable append + await controller.press_key("123") + + success = await controller.type_text(text, clear_existing=False) + return "Success" if success else "Failed" + + +@mcp.tool() +async def focus_and_clear_text( + ctx: Context, + coordinates: list[int], +) -> str: + """Focuses on a UI element at coordinates and clears its text content. + + 'coordinates' is a list [x, y] to tap first to gain focus. + """ + try: + controller = _get_controller() + except Exception as e: + return f"Error: Controller lazy initialization failed: {e}" + + if len(coordinates) != 2: + return "Error: coordinates must be [x, y]" + + err = await _ensure_focus_at_coords(controller, coordinates[0], coordinates[1]) + if err: + return f"Error focusing element: {err}" + + success = await controller.erase_text() + return "Success" if success else "Failed" + + +@mcp.tool() +async def erase_one_char(ctx: Context) -> str: + """Erases a single character (simulates Backspace).""" + try: + controller = _get_controller() + except Exception as e: + return f"Error: Controller lazy initialization failed: {e}" + + success = await controller.erase_text(nb_chars=1) + return "Success" if success else "Failed" + + +@mcp.tool() +async def press_key(ctx: Context, keycode: str) -> str: + """Presses a specific Android key event (e.g., KEYCODE_ENTER, KEYCODE_HOME).""" + try: + controller = _get_controller() + except Exception as e: + return f"Error: Controller lazy initialization failed: {e}" + + try: + success = await controller.press_key(keycode) + return "Success" if success else "Failed" + except Exception as e: + return f"Error: {e}" + + +@mcp.tool() +async def take_screenshot(ctx: Context) -> str: + """Takes a screenshot of the device screen. + + Returns the screenshot as a base64 encoded JPEG string. + """ + try: + controller = _get_controller() + except Exception as e: + return f"Error: Controller lazy initialization failed: {e}" + + try: + screenshot_b64 = await controller.take_screenshot() + return screenshot_b64 + except Exception as e: + return f"Error: {e}" + + +@mcp.tool() +async def get_ui_hierarchy(ctx: Context) -> str: + """Retrieves the current UI elements hierarchy from the device.""" + try: + controller = _get_controller() + except Exception as e: + logger.exception("Failed to initialize controller in get_ui_hierarchy") + raise e + + try: + elements = await controller.get_ui_elements() + return json.dumps(elements) + except Exception as e: + logger.exception("Failed to get UI elements in get_ui_hierarchy") + raise e + + +if __name__ == "__main__": + mcp.run(transport="stdio") diff --git a/artemis/mcp/xml_search_server.py b/artemis/mcp/xml_search_server.py new file mode 100644 index 0000000..70c25ff --- /dev/null +++ b/artemis/mcp/xml_search_server.py @@ -0,0 +1,280 @@ +# 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 os +from pathlib import Path +import re +import sys + +# Ensure repository root is in sys.path when executed directly or via MCP runner +_REPO_ROOT = str(Path(__file__).resolve().parent.parent.parent) +if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) + +# pylint: disable=wrong-import-position +from mcp.server.fastmcp import FastMCP + +from artemis.data_engine.storage import StorageManager + +mcp = FastMCP("Android_XML_Fuzzy_Search") + + +def get_subtree_info(element): + """Calculate total nodes and max depth of a subtree.""" + total_nodes = sum(1 for _ in element.iter()) - 1 + if total_nodes <= 0: + return 0, 0 + + def get_depth(el): + children = list(el) + if not children: + return 0 + return 1 + max(get_depth(c) for c in children) + + return total_nodes, get_depth(element) + + +def serialize_node(node, max_depth=1): + """Serialize an XML node to string with depth limit and truncation info.""" + + def _serialize(el, current_depth): + # Reconstruct attributes string + attrs = " ".join([f'{k}="{v}"' for k, v in el.attrib.items()]) + tag = el.tag + res = f"<{tag} {attrs}" if attrs else f"<{tag}" + + text = el.text.strip() if el.text else "" + children = list(el) + + if current_depth >= max_depth and children: + res += ">" + if text: + res += text + total_nodes, subtree_depth = get_subtree_info(el) + res += ( + f"\n\n" + ) + res += f"" + elif children or text: + res += ">" + if text: + res += text + for child in children: + res += "\n" + _serialize(child, current_depth + 1) + res += f"\n" + else: + res += " />" + return res + + return _serialize(node, 0) + + +@mcp.tool() +def search_ui(image_hash: str, query: str, threshold: float = 0.6) -> dict: + """Search for text in UI data (XML and OCR) retrieved from Data Engine by image hash. + + Args: + image_hash: The SHA-256 hash of the image. + query: The text to search for. + threshold: The similarity threshold (0.0 to 1.0). Defaults to 0.6. + + Returns: + A dictionary containing raw matches with bounds. + """ + db_path = os.getenv("DATA_ENGINE_DB_PATH") + if not db_path: + raise ValueError("DATA_ENGINE_DB_PATH environment variable is not set.") + base_dir = Path(db_path).parent + + try: + storage = StorageManager(db_path, base_dir) + except Exception as e: + return { + "error": f"Error initializing StorageManager: {e}", + "matches": [], + } + + matches = storage.search_ui_by_hash(image_hash, query, threshold) + + if not matches: + return {"matches": []} + + top_matches = matches[:8] + + results = [] + for m in top_matches: + # Extract bounds + box = None + if m["type"] == "xml": + bounds_str = m["node"].get("bounds") + if bounds_str: + box = parse_bounds(bounds_str) + elif m["type"] == "ocr": + position = m["node"].get("position", []) + if position: + left = min(v.get("x", 0) for v in position) + top = min(v.get("y", 0) for v in position) + right = max(v.get("x", 0) for v in position) + bottom = max(v.get("y", 0) for v in position) + box = (left, top, right, bottom) + + results.append( + { + "matched_text": m["matched_text"], + "type": m["type"], + "bounds": list(box) if box else None, + } + ) + + return {"matches": results} + + +def parse_bounds(bounds_str): + """Parse bounds string like '[left,top][right,bottom]'.""" + if not bounds_str: + return None + match = re.match(r"\[(\-?\d+),(\-?\d+)\]\[(\-?\d+),(\-?\d+)\]", bounds_str) + if match: + return tuple(map(int, match.groups())) + return None + + +@mcp.tool() +def search_by_coordinates(image_hash: str, x: int, y: int) -> str: + """Search for elements in UI data (XML and OCR) retrieved from Data Engine by image hash that overlap with the given coordinates. + + Args: + image_hash: The SHA-256 hash of the image. + x: The X coordinate. + y: The Y coordinate. + + Returns: + A string containing the matching nodes and OCR results. + """ + db_path = os.getenv("DATA_ENGINE_DB_PATH") + if not db_path: + raise ValueError("DATA_ENGINE_DB_PATH environment variable is not set.") + base_dir = Path(db_path).parent + + try: + storage = StorageManager(db_path, base_dir) + except Exception as e: + return f"Error initializing StorageManager: {e}" + + record = storage.get_image(image_hash) + if not record: + return "No record found for image hash." + + raw_xml_matches = [] + if record.ui_tree: + for node in record.ui_tree: + bounds_str = node.get("bounds") + if bounds_str: + bounds = parse_bounds(bounds_str) + if bounds: + left, top, right, bottom = bounds + if left <= x <= right and top <= y <= bottom: + raw_xml_matches.append(node) + + # Sort XML matches by area (smallest first) to prioritize "deepest/bottom" layers + def get_area(node): + bounds_str = node.get("bounds") + if bounds_str: + bounds = parse_bounds(bounds_str) + if bounds: + left, top, right, bottom = bounds + return (right - left) * (bottom - top) + return float("inf") + + raw_xml_matches.sort(key=get_area) + + # Collect bottom layers up to the first interactable node (max 2 nodes) + interaction_keys = { + "clickable", + "scrollable", + "long-clickable", + "checkable", + "focusable", + "editable", + "selected", + } + xml_matches = [] + for node in raw_xml_matches: + has_text = bool(node.get("text")) or bool(node.get("content-desc")) + is_interactable = any(node.get(k) == "true" for k in interaction_keys) + + # Skip completely empty/useless middle wrappers + if not has_text and not is_interactable: + continue + + xml_matches.append(node) + + # Stop collecting if we hit an interactable node, or we reached the max of 2 nodes + if is_interactable or len(xml_matches) >= 2: + break + + ocr_matches = [] + if record.ocr_result: + for ocr in record.ocr_result: + position = ocr.get("position", []) + if position: + left = min(v.get("x", 0) for v in position) + top = min(v.get("y", 0) for v in position) + right = max(v.get("x", 0) for v in position) + bottom = max(v.get("y", 0) for v in position) + + if left <= x <= right and top <= y <= bottom: + ocr_matches.append(ocr) + + if not xml_matches and not ocr_matches: + return f"No elements found at coordinates ({x}, {y})." + + output = [] + output.append(f"Found matches overlapping with ({x}, {y}):") + + if xml_matches: + for i, node in enumerate(xml_matches, 1): + allowed_keys = { + "class", + "text", + "resource-id", + "content-desc", + "clickable", + "enabled", + "checked", + "focused", + "scrollable", + } + # Drop empty strings and None, but keep explicit "false" or "0" + node_copy = {k: v for k, v in node.items() if k in allowed_keys and v not in ("", None)} + + # Extract class to put it first, then format the rest + cls_name = node_copy.pop("class", "UnknownClass") + props = [f"{k}: {v}" for k, v in node_copy.items()] + + line = f"Node {i}: {cls_name}" + if props: + line += " | " + " | ".join(props) + output.append(line) + + if ocr_matches: + for i, ocr in enumerate(ocr_matches[:5], 1): # Limit to top 5 + output.append(f"OCR {i}: Text: {ocr.get('text')}") + + return " ".join(output) + + +if __name__ == "__main__": + mcp.run() diff --git a/artemis/sdk/__init__.py b/artemis/sdk/__init__.py new file mode 100644 index 0000000..414a22a --- /dev/null +++ b/artemis/sdk/__init__.py @@ -0,0 +1,26 @@ +# 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. + +"""ARTEMIS SDK for running mobile automation tasks. + +This package provides APIs for interacting with mobile devices and executing +tasks. +""" + +from artemis.sdk import builders, types +from artemis.sdk.agent import Agent + +__all__ = ["Agent"] +__all__ += types.__all__ +__all__ += builders.__all__ diff --git a/artemis/sdk/agent.py b/artemis/sdk/agent.py new file mode 100644 index 0000000..4a4b4f0 --- /dev/null +++ b/artemis/sdk/agent.py @@ -0,0 +1,1028 @@ +# 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 asyncio +import os + +try: + from datetime import UTC, datetime +except ImportError: + from datetime import datetime + + UTC = UTC +from pathlib import Path +from platform import system +import shutil +from shutil import which +import sys +from types import NoneType +from typing import TypeVar, overload +import uuid + +from adbutils import AdbClient +from dotenv import load_dotenv +from google import genai +from langchain_core.messages import AIMessage +from langchain_google_genai import ChatGoogleGenerativeAI +from PIL import Image +from pydantic import BaseModel + +from artemis.agents.flash.runner import FlashRunner +from artemis.agents.outputter.outputter import outputter +from artemis.clients.ui_automator_client import UIAutomatorClient +from artemis.config import OutputConfig, record_events, settings +from artemis.constants import RECURSION_LIMIT +from artemis.context import ( + ArtemisContext, + DeviceContext, + DevicePlatform, + ExecutionSetup, +) +from artemis.controllers.controller_factory import get_controller +from artemis.controllers.platform_specific_commands_controller import ( + get_first_device, +) +from artemis.data_engine.engine import DataEngine +from artemis.data_engine.trace import DataEngineCallbackHandler +from artemis.graph.graph import get_graph +from artemis.graph.state import State +from artemis.sdk.builders.agent_config_builder import get_default_agent_config +from artemis.sdk.builders.task_request_builder import TaskRequestBuilder +from artemis.sdk.types.agent import AgentConfig +from artemis.sdk.types.exceptions import ( + AgentError, + AgentNotInitializedError, + AgentProfileNotFoundError, + AgentTaskRequestError, + DeviceNotFoundError, + ExecutableNotFoundError, +) +from artemis.sdk.types.task import ( + AgentProfile, + Task, + TaskRequest, +) +from artemis.utils.app_launch_utils import _handle_initial_app_launch +from artemis.utils.logger import get_logger +from artemis.utils.media import ( + create_gif_from_trace_folder, + create_steps_json_from_trace_folder, + remove_images_from_trace_folder, + remove_steps_json_from_trace_folder, +) + +logger = get_logger(__name__) + +TOutput = TypeVar("TOutput", bound=BaseModel | None) + +load_dotenv() + + +class Agent: + _config: AgentConfig + _tasks: list[Task] = [] + _tmp_traces_dir: Path + _initialized: bool = False + _device_context: DeviceContext + _adb_client: AdbClient | None + _ui_adb_client: UIAutomatorClient | None + + _current_task: asyncio.Task | None = None + _task_lock: asyncio.Lock + _cloud_mobile_id: str | None = None + + def __init__(self, *, config: AgentConfig | None = None): + self._config = config or get_default_agent_config() + self._tasks = [] + self._tmp_traces_dir = Path(settings.TRACES_PATH) + self._initialized = False + self._task_lock = asyncio.Lock() + + async def init( + self, + api_key: str | None = None, + retry_count: int = 5, + retry_wait_seconds: int = 5, + ): + + return await self._init_internal( + api_key=api_key, + retry_count=retry_count, + retry_wait_seconds=retry_wait_seconds, + ) + + async def _init_internal( + self, + api_key: str | None = None, + retry_count: int = 5, + retry_wait_seconds: int = 5, + ): + + if os.environ.get("ARTEMIS_CLOUD_MODE") != "1" and not which("adb"): + raise ExecutableNotFoundError("adb") + + if self._initialized: + logger.warning("Agent is already initialized. Skipping...") + return True + + # Get first available device ID + if os.environ.get("ARTEMIS_CLOUD_MODE") == "1": + device_id = os.environ.get("ADB_DEVICE_SERIAL", "cloud_device") + platform = DevicePlatform.ANDROID + elif not self._config.device_id or not self._config.device_platform: + device_id, platform, _ = get_first_device(logger=logger) + else: + device_id, platform = ( + self._config.device_id, + self._config.device_platform, + ) + + if not device_id or not platform: + error_msg = "No device found. Exiting." + logger.error(error_msg) + raise DeviceNotFoundError(error_msg) + + # Initialize clients + if os.environ.get("ARTEMIS_CLOUD_MODE") != "1": + self._init_clients( + device_id=device_id, + platform=platform, + ) + else: + from cloud_service.virtualization import RemoteAdbClient, RemoteUIAutomatorClient + + self._adb_client = RemoteAdbClient() + self._ui_adb_client = RemoteUIAutomatorClient( + device_id=device_id, + adb_client=self._adb_client, + ) + + self._device_context = await self._get_device_context( + device_id=device_id, platform=platform + ) + logger.info(self._device_context.to_str()) + + # Asynchronously pre-warm LLM connection pools in the background + asyncio.create_task(self._prewarm_llm_connections(api_key)) + + logger.info("โœ… Artemis agent initialized.") + self._initialized = True + + return True + + async def _prewarm_llm_connections(self, api_key: str | None = None): + """Pre-warms the HTTP2/gRPC connection pools for both Native GenAI and LangChain clients in the background.""" + logger.info("Starting background pre-warming of Gemini API connection pools...") + try: + key = api_key + if not key and settings.GOOGLE_API_KEY: + key = settings.GOOGLE_API_KEY.get_secret_value() + + if not key: + logger.warning("Skipping LLM pre-warming: No API key available.") + return + + # 1. Pre-warm Native SDK client + client = genai.Client(api_key=key) + + # 2. Pre-warm LangChain client + chat = ChatGoogleGenerativeAI(model="gemini-3.6-flash", google_api_key=key) + + # Fire both calls concurrently in the background + await asyncio.gather( + client.aio.models.count_tokens(model="gemini-3.6-flash", contents="ping"), + chat.ainvoke("ping"), + return_exceptions=True, + ) + logger.success("Gemini API connection pools successfully pre-warmed.") + except Exception as e: + logger.warning(f"Failed to pre-warm LLM connections: {e}") + + async def install_apk(self, apk_path: str | Path) -> None: + """Install an APK on the connected device. + + For cloud mobiles, the APK must be x86_64 compatible. + + Args: + apk_path: Path to the local APK file to install + + Raises: + AgentNotInitializedError: If the agent is not initialized + AgentError: If attempting to install on non-Android device or ADB + operations fail + FileNotFoundError: If the APK file doesn't exist + CloudMobileServiceUninitializedError: If cloud service is + unavailable + """ + await self._install_apk_internal(apk_path) + + async def _install_apk_internal(self, apk_path: str | Path) -> None: + if isinstance(apk_path, str): + apk_path = Path(apk_path) + + if not apk_path.exists(): + raise FileNotFoundError(f"APK file not found: {apk_path}") + + if not self._initialized: + raise AgentNotInitializedError() + + device_id = self._device_context.device_id + logger.info(f"Installing APK on Android device '{device_id}'") + if not self._adb_client: + raise AgentError("ADB client not initialized") + + device = self._adb_client.device(serial=device_id) + await asyncio.to_thread(device.install, apk_path) + logger.info(f"APK installed successfully on Android device '{device_id}'") + + async def install_app(self, app_path: str | Path) -> str | None: + """Install an app on the connected device. + + For Android: Installs an APK file using ADB. + + Args: + app_path: Path to the app to install (Android APK file). + + Returns: + None. + + Raises: + AgentNotInitializedError: If the agent is not initialized + AgentError: If installation fails or platform is unsupported + FileNotFoundError: If the app file/folder doesn't exist + """ + return await self._install_app_internal(app_path) + + async def _install_app_internal(self, app_path: str | Path) -> str | None: + if isinstance(app_path, str): + app_path = Path(app_path) + + if not app_path.exists(): + raise FileNotFoundError(f"App not found: {app_path}") + + if not self._initialized: + raise AgentNotInitializedError() + + await self._install_apk_internal(app_path) + return None + + def new_task(self, goal: str): + """Create a new task request builder. + + Args: + goal: Natural language description of what to accomplish + + Returns: + TaskRequestBuilder that can be configured with: + - .with_output_format() for structured output + - .with_output_description() for output description + - .with_locked_app_package() to restrict execution to a specific app + - .using_profile() to specify an LLM profile + - .with_max_steps() to set maximum execution steps + - .with_trace_recording() to enable trace recording + - .with_name() to set a custom task name + """ + return TaskRequestBuilder[None].from_common( + goal=goal, + common=self._config.task_request_defaults, + ) + + @overload + async def run_task( + self, + *, + goal: str, + output: type[TOutput], + profile: str | AgentProfile | None = None, + name: str | None = None, + locked_app_package: str | None = None, + app_path: str | Path | None = None, + ) -> TOutput | None: ... + + @overload + async def run_task( + self, + *, + goal: str, + output: str, + profile: str | AgentProfile | None = None, + name: str | None = None, + locked_app_package: str | None = None, + app_path: str | Path | None = None, + ) -> str | dict | None: ... + + @overload + async def run_task( + self, + *, + goal: str, + output=None, + profile: str | AgentProfile | None = None, + name: str | None = None, + locked_app_package: str | None = None, + app_path: str | Path | None = None, + ) -> str | None: ... + + @overload + async def run_task( + self, + *, + request: TaskRequest[None], + locked_app_package: str | None = None, + app_path: str | Path | None = None, + ) -> str | dict | None: ... + + @overload + async def run_task( + self, + *, + request: TaskRequest[TOutput], + locked_app_package: str | None = None, + app_path: str | Path | None = None, + ) -> TOutput | None: ... + + async def run_task( + self, + *, + goal: str | None = None, + output: type[TOutput] | str | None = None, + profile: str | AgentProfile | None = None, + locked_app_package: str | None = None, + name: str | None = None, + app_path: str | Path | None = None, + request: TaskRequest[TOutput] | None = None, + ) -> str | dict | TOutput | None: + + # Normal local execution path + if request is not None: + if locked_app_package is not None: + if request.locked_app_package: + logger.warning( + "Locked app package specified both in the request and" + " as a parameter. Using the parameter value." + ) + request.locked_app_package = locked_app_package + # Handle app_path parameter override + if app_path is not None: + if request.app_path: + logger.warning( + "App path specified both in the request and as a" + " parameter. Using the parameter value." + ) + request.app_path = Path(app_path) if isinstance(app_path, str) else app_path + return await self._run_task(request=request) + if goal is None: + raise AgentTaskRequestError("Goal is required") + task_request = self.new_task(goal=goal) + if output is not None: + if isinstance(output, str): + task_request.with_output_description(description=output) + elif output is not NoneType: + task_request.with_output_format(output_format=output) + if profile is not None: + task_request.using_profile(profile=profile) + if name is not None: + task_request.with_name(name=name) + if locked_app_package is not None: + task_request.with_locked_app_package(package_name=locked_app_package) + if app_path is not None: + task_request.with_app_path(app_path=app_path) + return await self._run_task(task_request.build()) + + async def _run_task( + self, + request: TaskRequest[TOutput], + ) -> str | dict | TOutput | None: + if not self._initialized: + raise AgentNotInitializedError() + + if request.profile: + agent_profile = self._config.agent_profiles.get(request.profile) + if agent_profile is None: + if request.profile.lower() in ( + "flash", + "pro", + "ultra", + "default", + ): + agent_profile = self._config.default_profile + else: + raise AgentProfileNotFoundError(request.profile) + else: + agent_profile = self._config.default_profile + + logger.info(str(agent_profile)) + + on_status_changed = None + task_id = str(uuid.uuid4()) + + task = Task( + id=task_id, + device=self._device_context, + status="pending", + request=request, + created_at=datetime.now(), + on_status_changed=on_status_changed, + ) + self._tasks.append(task) + task_name = task.get_name() + + context = ArtemisContext( + device=self._device_context, + adb_client=self._adb_client, + ui_adb_client=self._ui_adb_client, + llm_config=agent_profile.llm_config, + agent_config=self._config, + ) + + self._prepare_tracing(task=task, context=context) + await self._prepare_app_installation(task=task) + await self._prepare_device_environment(context=context) + await self._prepare_app_lock(task=task, context=context) + self._prepare_output_files(task=task) + + output_config = None + if request.output_description or request.output_format: + output_config = OutputConfig( + output_description=request.output_description, + structured_output=request.output_format, # type: ignore + ) + logger.info(str(output_config)) + + logger.info(f"[{task_name}] Starting graph with goal: `{request.goal}`") + state = self._get_graph_state(task=task) + graph_input = state.model_dump() + datetime.now(UTC) + + async def _execute_task_logic(): + last_state: State | None = None + last_state_snapshot: dict | None = None + output = None + try: + async with context: + recording_started = False + if self._config.video_recording_tools_enabled: + try: + controller = get_controller(context) + + output_dir = None + if context.execution_setup and context.execution_setup.traces_path: + output_dir = ( + self._tmp_traces_dir / context.execution_setup.trace_name + ).resolve() + + logger.info(f"[{task_name}] Starting automated screen recording...") + start_res = await controller.start_video_recording( + output_dir=output_dir + ) + if start_res and start_res.success: + recording_started = True + else: + logger.warning( + f"[{task_name}] Failed to start screen" + f" recording: {start_res.message if start_res else 'unknown'}" + ) + except Exception as e: + logger.error(f"[{task_name}] Failed to start screen recording: {e}") + + try: + if request.profile and request.profile.lower() == "flash": + logger.info(f"[{task_name}] Invoking FlashRunner reactive loop...") + await task.set_status( + status="running", + message="Invoking FlashRunner...", + ) + + max_turns = ( + request.max_steps if request.max_steps != RECURSION_LIMIT else 35 + ) + runner = FlashRunner(context, goal=request.goal, max_turns=max_turns) + flash_result = await runner.run(state) + + output = flash_result + last_state_snapshot = state.model_dump() + + status = flash_result.get("status") + if status == "completed": + logger.info(f"โœ… Automation '{task_name}' is success โœ…") + await task.finalize(content=output, state=last_state_snapshot) + if context.data_engine: + context.data_engine.end_session("success") + else: + err = ( + f"[{task_name}] FlashRunner failed:" + f" {flash_result.get('explanation')}" + ) + logger.warning(err) + await task.finalize( + content=output, + state=last_state_snapshot, + error=err, + ) + if context.data_engine: + context.data_engine.end_session("failed") + return output + else: + logger.info(f"[{task_name}] Invoking graph with input: {graph_input}") + await task.set_status(status="running", message="Invoking graph...") + async for chunk in (await get_graph(context)).astream( + input=graph_input, + config={ + "recursion_limit": task.request.max_steps, + "callbacks": ( + ( + list(self._config.graph_config_callbacks) + if self._config.graph_config_callbacks is not None + else [] + ) + + [DataEngineCallbackHandler(context)] + if context.data_engine + else self._config.graph_config_callbacks + ), + }, + stream_mode=[ + "messages", + "custom", + "updates", + "values", + ], + ): + stream_mode, payload = chunk + if stream_mode == "values": + last_state_snapshot = payload # type: ignore + last_state = State(**last_state_snapshot) # type: ignore + + if not last_state: + err = f"[{task_name}] No result received from graph" + logger.warning(err) + await task.finalize( + content=output, + state=last_state_snapshot, + error=err, + ) + if context.data_engine: + context.data_engine.end_session("failed") + return None + + print_ai_response_to_stderr(graph_result=last_state) + output = await self._extract_output( + task_name=task_name, + ctx=context, + request=request, + output_config=output_config, + state=last_state, + ) + logger.info(f"โœ… Automation '{task_name}' is success โœ…") + await task.finalize(content=output, state=last_state_snapshot) + if context.data_engine: + context.data_engine.end_session("success") + + return output + finally: + if recording_started: + try: + logger.info(f"[{task_name}] Stopping automated screen recording...") + + controller = get_controller(context) + stop_result = await controller.stop_video_recording() + if stop_result.success and stop_result.video_path: + logger.info( + f"[{task_name}] Screen recording saved" + f" to: {stop_result.video_path}" + ) + else: + logger.warning( + f"[{task_name}] Failed to save screen" + f" recording: {stop_result.message}" + ) + except Exception as e: + logger.error(f"[{task_name}] Error stopping screen recording: {e}") + except asyncio.CancelledError: + err = f"[{task_name}] Task cancelled" + logger.warning(err) + await task.finalize( + content=output, + state=last_state_snapshot, + error=err, + cancelled=True, + ) + if context.data_engine: + context.data_engine.end_session("cancelled") + + raise + except Exception as e: + err = f"[{task_name}] Error running automation: {e}" + logger.error(err) + await task.finalize( + content=output, + state=last_state_snapshot, + error=err, + ) + if context.data_engine: + context.data_engine.end_session("failed") + + raise + finally: + await self._finalize_tracing(task=task, context=context) + + async with self._task_lock: + if self._current_task and not self._current_task.done(): + logger.warning( + "Another automation task is already running. " + "Stopping it before starting the new one." + ) + self.stop_current_task() + try: + await self._current_task + except asyncio.CancelledError: + pass + + try: + self._current_task = asyncio.create_task(_execute_task_logic()) + return await self._current_task + finally: + self._current_task = None + + def stop_current_task(self): + """Requests cancellation of the currently running automation task.""" + if self._current_task and not self._current_task.done(): + logger.info("Requesting to stop the current automation task...") + was_cancelled = self._current_task.cancel() + if was_cancelled: + logger.success("Cancellation request for the current task was sent.") + else: + logger.warning( + "Could not send cancellation request for the current task " + "(it may already be completing)." + ) + else: + logger.info("No active automation task to stop.") + + async def get_screenshot(self) -> Image.Image: + """Capture a screenshot from the mobile device. + + Returns: + Screenshot as PIL Image + + Raises: + AgentNotInitializedError: If the agent is not initialized + Exception: If screenshot capture fails + """ + if not self._initialized: + raise AgentNotInitializedError() + + # Use ADB to capture screenshot + logger.info("Capturing screenshot from local Android device") + if not self._adb_client: + raise Exception("ADB client not initialized") + + device = self._adb_client.device(serial=self._device_context.device_id) + screenshot = await asyncio.to_thread(device.screenshot) + logger.info("Screenshot captured from local Android device") + return screenshot + + async def clean(self, force: bool = False): + if not self._initialized and not force: + return + + self._initialized = False + logger.info("โœ… Artemis agent stopped.") + + async def _prepare_app_installation(self, task: Task) -> str | None: + """Install app if app_path is specified in the task request. + + Returns: + None. + """ + if not task.request.app_path: + return None + + task_name = task.get_name() + logger.info(f"[{task_name}] Installing app from: {task.request.app_path}") + + await self.install_app(task.request.app_path) + + return None + + async def _prepare_device_environment(self, context: ArtemisContext): + """Prepare device environment flags (like forcing Web Accessibility) before the task runs.""" + if not self._config.force_web_accessibility: + logger.info( + "Forcing web accessibility is disabled in AgentConfig. Skipping" + " device environment prep..." + ) + return + + if not self._adb_client or not self._device_context: + logger.warning( + "ADB client or device context not available. Skipping device environment prep..." + ) + return + + device_id = self._device_context.device_id + try: + device = self._adb_client.device(serial=device_id) + logger.info(f"[{device_id}] โš™๏ธ Configuring Chrome/WebView Web Accessibility flags...") + + # 1. Write the force accessibility command flag to Chrome command line + await asyncio.to_thread( + device.shell, + "echo 'chrome --force-renderer-accessibility' >" + " /data/local/tmp/chrome-command-line", + ) + await asyncio.to_thread(device.shell, "chmod 555 /data/local/tmp/chrome-command-line") + + # 2. Write the same command line flag to System WebView config file + await asyncio.to_thread( + device.shell, + "echo 'chrome --force-renderer-accessibility' >" + " /data/local/tmp/webview-command-line", + ) + await asyncio.to_thread(device.shell, "chmod 555 /data/local/tmp/webview-command-line") + + # 3. Force stop Chrome so the flag takes effect next time it launches + await asyncio.to_thread(device.shell, "am force-stop com.android.chrome") + logger.success( + f"[{device_id}] โœ… Chrome Web Accessibility command line flags" + " configured successfully." + ) + except Exception as e: + logger.warning( + "Failed to configure Chrome Web Accessibility flags (needs" + f" root/userdebug device): {e}" + ) + + async def _prepare_app_lock(self, task: Task, context: ArtemisContext): + """Prepare app lock by launching the locked app if specified.""" + if not task.request.locked_app_package: + return + + task_name = task.get_name() + logger.info(f"[{task_name}] Preparing app lock for: {task.request.locked_app_package}") + + app_lock_status = await _handle_initial_app_launch( + ctx=context, locked_app_package=task.request.locked_app_package + ) + + if context.execution_setup is None: + context.execution_setup = ExecutionSetup(app_lock_status=app_lock_status) + else: + context.execution_setup.app_lock_status = app_lock_status + + if app_lock_status.locked_app_initial_launch_success is False: + error = app_lock_status.locked_app_initial_launch_error + logger.warning(f"[{task_name}] Failed to launch locked app: {error}") + + def _prepare_tracing(self, task: Task, context: ArtemisContext): + """Prepare tracing and data engine setup.""" + task_name = task.get_name() + temp_trace_path = Path(self._tmp_traces_dir / task_name).resolve() + temp_trace_path.mkdir(parents=True, exist_ok=True) + + if task.request.record_trace: + traces_output_path = Path(task.request.trace_path).resolve() + logger.info(f"[{task_name}] ๐Ÿ“‚ Traces output path: {traces_output_path}") + logger.info(f"[{task_name}] ๐Ÿ“„๐Ÿ“‚ Traces temp path: {temp_trace_path}") + traces_output_path.mkdir(parents=True, exist_ok=True) + + context.execution_setup = ExecutionSetup( + traces_path=self._tmp_traces_dir, + trace_name=task_name, + enable_remote_tracing=task.request.enable_remote_tracing + if task.request.record_trace + else False, + video_recording_tools_enabled=self._config.video_recording_tools_enabled, + disable_checker=self._config.disable_checker, + checker_max_iterations=self._config.checker_max_iterations, + checker_max_chat_rounds=self._config.checker_max_chat_rounds, + disable_planner_validation=self._config.disable_planner_validation, + planner_validation_threshold=self._config.planner_validation_threshold, + enable_committee=self._config.enable_committee, + committee_debate_rounds=self._config.committee_debate_rounds, + disable_outputter=self._config.disable_outputter, + outputter=self._config.outputter, + explorer=self._config.explorer, + explorer_versions=self._config.explorer_versions, + ) + + context.data_engine = DataEngine(ctx=context) + device_data = context.device.model_dump() if context.device else {} + if task.request.profile: + device_data["profile"] = task.request.profile + context.data_engine.start_session(goal=task.request.goal, device_info=device_data) + + async def _finalize_tracing(self, task: Task, context: ArtemisContext): + if context.data_engine: + session_status = "completed" + if task.status == "failed": + session_status = "failed" + elif task.status == "cancelled": + session_status = "cancelled" + context.data_engine.end_session(status=session_status) + + exec_setup_ctx = context.execution_setup + if not exec_setup_ctx: + return + + if exec_setup_ctx.traces_path is None or exec_setup_ctx.trace_name is None: + return + + task_name = task.get_name() + temp_trace_path = (self._tmp_traces_dir / exec_setup_ctx.trace_name).resolve() + + if not task.request.record_trace: + try: + if temp_trace_path.exists(): + shutil.rmtree(temp_trace_path) + logger.info(f"[{task_name}] Cleaned up temporary trace folder.") + except Exception as e: + logger.error(f"[{task_name}] Failed to clean up temp trace folder: {e}") + return + + status = "_PASS" if task.status == "completed" else "_FAIL" + ts = task.created_at.strftime("%Y-%m-%dT%H-%M-%S") + new_name = f"{exec_setup_ctx.trace_name}{status}_{ts}" + + traces_output_path = Path(task.request.trace_path).resolve() + + logger.info(f"[{task_name}] Compiling trace FROM FOLDER: " + str(temp_trace_path)) + create_gif_from_trace_folder(temp_trace_path) + create_steps_json_from_trace_folder(temp_trace_path) + + logger.info(f"[{task_name}] Video created, removing dust...") + remove_images_from_trace_folder(temp_trace_path) + remove_steps_json_from_trace_folder(temp_trace_path) + logger.info(f"[{task_name}] ๐Ÿ“ฝ๏ธ Trace compiled, moving to output path ๐Ÿ“ฝ๏ธ") + + output_folder_path = temp_trace_path.rename(traces_output_path / new_name).resolve() + logger.info(f"[{task_name}] ๐Ÿ“‚โœ… Traces located in: {output_folder_path}") + + new_video_path = output_folder_path / "recording.mp4" + if new_video_path.exists() and context.data_engine: + context.data_engine.update_video_path(new_video_path) + + def _prepare_output_files(self, task: Task): + if task.request.llm_output_path: + _validate_and_prepare_file(file_path=task.request.llm_output_path) + + async def _extract_output( + self, + task_name: str, + ctx: ArtemisContext, + request: TaskRequest[TOutput], + output_config: OutputConfig | None, + state: State, + ) -> str | dict | TOutput | None: + exec_setup = getattr(ctx, "execution_setup", None) + outputter_cfg = getattr(exec_setup, "outputter", None) or ( + getattr(ctx.agent_config, "outputter", None) if ctx and ctx.agent_config else None + ) + outputter_enabled = getattr(outputter_cfg, "enabled", True) if outputter_cfg else True + force_synthesis = ( + getattr(outputter_cfg, "force_synthesis", False) if outputter_cfg else False + ) + if exec_setup and getattr(exec_setup, "disable_outputter", False): + outputter_enabled = False + + should_run = outputter_enabled and ( + ( + output_config + and output_config.needs_structured_format(default_enabled=outputter_enabled) + ) + or force_synthesis + or (output_config and output_config.enable_outputter is True) + ) + if output_config and output_config.enable_outputter is False: + should_run = False + + if should_run: + logger.info(f"[{task_name}] Generating structured output via Outputter...") + try: + structured_output = await outputter( + ctx=ctx, + output_config=output_config or OutputConfig(), + graph_output=state, + ) + logger.info(f"[{task_name}] Structured output: {structured_output}") + record_events( + output_path=request.llm_output_path, + events=structured_output, + ) + if request.output_format is not None and request.output_format is not NoneType: + return request.output_format.model_validate(structured_output) + return structured_output + except Exception as e: + logger.error(f"[{task_name}] Failed to generate structured output: {e}") + return None + if state and state.messages: + for msg in reversed(state.messages): + if msg.content: + logger.info(f"Fallback output from message: {msg.content}") + record_events(output_path=request.llm_output_path, events=msg.content) + return msg.content + return None + + def _get_graph_state(self, task: Task): + return State( + messages=[], + initial_goal=task.request.goal, + subgoal_plan=[], + latest_ui_hierarchy=None, + latest_screenshot=None, + focused_app_info=None, + device_date=None, + structured_decisions=None, + complete_subgoals_by_ids=[], + remaining_steps=task.request.max_steps, + validator_messages=[], + ) + + def _init_clients( + self, + device_id: str, + platform: DevicePlatform, + ): + self._adb_client = AdbClient( + host=self._config.servers.adb_host, + port=self._config.servers.adb_port, + ) + self._ui_adb_client = UIAutomatorClient(device_id=device_id) + + async def _get_device_context( + self, + device_id: str, + platform: DevicePlatform, + ) -> DeviceContext: + + host_platform = system() + if os.environ.get("ARTEMIS_CLOUD_MODE") == "1": + width, height = (1080, 2400) + if self._adb_client: + try: + width, height = self._adb_client.device(device_id).window_size() + logger.info( + f"Retrieved remote device screen dimensions dynamically: {width}x{height}" + ) + except Exception as e: + logger.debug( + f"Failed to query remote window size: {e}. Defaulting to 1080x2400." + ) + return DeviceContext( + host_platform="LINUX", + mobile_platform=platform, + device_id=device_id, + device_width=width, + device_height=height, + ) + + # Get real device dimensions from the device + if self._ui_adb_client: + try: + # Use UIAutomator to get actual screen dimensions + screen_data = self._ui_adb_client.get_screen_data() + device_width = screen_data.width + device_height = screen_data.height + logger.info(f"Retrieved Android screen dimensions: {device_width}x{device_height}") + except Exception as e: + logger.warning(f"Failed to get Android screen dimensions: {e}, using defaults") + device_width, device_height = 1080, 2340 + else: + logger.warning("UIAutomator client not available, using default dimensions") + device_width, device_height = 1080, 2340 + + return DeviceContext( + host_platform="WINDOWS" if host_platform == "Windows" else "LINUX", + mobile_platform=platform, + device_id=device_id, + device_width=device_width, + device_height=device_height, + ) + + +def _validate_and_prepare_file(file_path: Path): + path_obj = Path(file_path) + if path_obj.exists() and path_obj.is_dir(): + raise AgentTaskRequestError(f"Error: Path '{file_path}' is a directory, not a file.") + try: + path_obj.parent.mkdir(parents=True, exist_ok=True) + path_obj.touch(exist_ok=True) + except OSError as e: + raise AgentTaskRequestError(f"Error creating file '{file_path}': {e}") + + +def print_ai_response_to_stderr(graph_result: State): + for msg in reversed(graph_result.messages): + if isinstance(msg, AIMessage): + print(msg.content, file=sys.stderr) + return diff --git a/artemis/sdk/builders/__init__.py b/artemis/sdk/builders/__init__.py new file mode 100644 index 0000000..dbdba77 --- /dev/null +++ b/artemis/sdk/builders/__init__.py @@ -0,0 +1,29 @@ +# 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. + +"""Builder classes for configuring ARTEMIS components.""" + +from artemis.sdk.builders.agent_config_builder import AgentConfigBuilder +from artemis.sdk.builders.index import Builders +from artemis.sdk.builders.task_request_builder import ( + TaskRequestBuilder, + TaskRequestCommonBuilder, +) + +__all__ = [ + "AgentConfigBuilder", + "TaskRequestCommonBuilder", + "TaskRequestBuilder", + "Builders", +] diff --git a/artemis/sdk/builders/agent_config_builder.py b/artemis/sdk/builders/agent_config_builder.py new file mode 100644 index 0000000..c3c3937 --- /dev/null +++ b/artemis/sdk/builders/agent_config_builder.py @@ -0,0 +1,528 @@ +# 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. + +"""Builder for AgentConfig objects using a fluent interface.""" + +import copy +import os +from typing import Any + +from langchain_core.callbacks.base import Callbacks +from artemis.config import ( + ExplorerVersion, + get_default_llm_config, + load_agent_config, + settings, +) +from artemis.context import DevicePlatform +from artemis.sdk.constants import DEFAULT_PROFILE_NAME +from artemis.sdk.types.agent import ( + AgentConfig, + AgentProfile, + ServerConfig, +) +from artemis.sdk.types.task import TaskRequestCommon +from artemis.utils.video import detect_video_tools_enabled + + +class AgentConfigBuilder: + """Builder class providing a fluent interface for creating AgentConfig objects. + + This builder allows for step-by-step construction of an AgentConfig with + clear methods that make the configuration process intuitive and type-safe. + + Examples: + >>> builder = AgentConfigBuilder() + >>> config = (builder + ... .add_profile(AgentProfile(name="HighReasoning", + llm_config=LLMConfig(...))) + ... .add_profile(AgentProfile(name="LowReasoning", + llm_config=LLMConfig(...))) + ... .for_device(DevicePlatform.ANDROID, "device123") + ... .with_default_task_config(TaskRequestCommon(max_steps=30)) + ... .with_default_profile("HighReasoning") + ... .build() + ... ) + """ + + def __init__(self): + """Initialize an empty AgentConfigBuilder.""" + self._agent_profiles: dict[str, AgentProfile] = {} + self._task_request_defaults: TaskRequestCommon | None = None + self._default_profile: str | AgentProfile | None = None + self._device_id: str | None = None + self._device_platform: DevicePlatform | None = None + self._servers: ServerConfig = get_default_servers() + self._graph_config_callbacks: Callbacks = None + self._video_recording_tools_enabled: bool = detect_video_tools_enabled() + self._force_web_accessibility: bool = False + self._disable_checker: bool = True + self._cloud_mobile_id_or_ref: str | None = None + + agent_cfg = load_agent_config() + self._explorer = agent_cfg.explorer + self._explorer_versions = agent_cfg.explorer_versions + self._blacklisted_tools = agent_cfg.blacklisted_tools + self._enable_video_ledger = agent_cfg.video_analyzer.enable_ledger + if agent_cfg.video_analyzer.enabled is not None: + self._video_recording_tools_enabled = agent_cfg.video_analyzer.enabled + else: + self._video_recording_tools_enabled = detect_video_tools_enabled() + self._disable_planner_validation = not agent_cfg.planner_validation.enabled + self._planner_validation_threshold = agent_cfg.planner_validation.similarity_threshold + self._enable_committee = agent_cfg.committee.enabled + self._committee_debate_rounds = agent_cfg.committee.debate_rounds + self._disable_checker = not agent_cfg.checker.enabled + self._checker_max_iterations = agent_cfg.checker.max_iterations + self._checker_max_chat_rounds = agent_cfg.checker.max_chat_rounds + self._outputter = agent_cfg.outputter + self._disable_outputter = not agent_cfg.outputter.enabled + self._flash = agent_cfg.flash + self._pro = agent_cfg.pro + + def add_profile(self, profile: AgentProfile, validate: bool = True) -> "AgentConfigBuilder": + """Add an agent profile to the ARTEMIS agent. + + Args: + profile: The agent profile to add + """ + self._agent_profiles[profile.name] = profile + if validate: + profile.llm_config.validate_providers() + return self + + def add_profiles( + self, + profiles: list[AgentProfile], + validate: bool = True, + ) -> "AgentConfigBuilder": + """Add multiple agent profiles to the ARTEMIS agent. + + Args: + profiles: List of agent profiles to add + """ + for profile in profiles: + self.add_profile(profile=profile, validate=validate) + return self + + def with_default_profile(self, profile: str | AgentProfile) -> "AgentConfigBuilder": + """Set the default agent profile used for tasks. + + Args: + profile: The name or instance of the default agent profile + """ + self._default_profile = profile + return self + + def for_device( + self, + platform: DevicePlatform, + device_id: str, + ) -> "AgentConfigBuilder": + """Configure the ARTEMIS agent for a specific device. + + Args: + platform: The device platform (ANDROID) + device_id: The unique identifier for the device + """ + if self._cloud_mobile_id_or_ref is not None: + raise ValueError( + "Device ID cannot be set when a cloud mobile is already" + " configured.\n> for_device() and for_cloud_mobile() are" + " mutually exclusive" + ) + self._device_id = device_id + self._device_platform = platform + return self + + def with_default_task_config(self, config: TaskRequestCommon) -> "AgentConfigBuilder": + """Set the default task configuration. + + Args: + config: The task configuration to use as default + """ + self._task_request_defaults = copy.deepcopy(config) + return self + + def with_adb_server(self, host: str, port: int | None = None) -> "AgentConfigBuilder": + """Set the ADB server host and port. + + Args: + host: The ADB server host + port: The ADB server port + """ + self._servers.adb_host = host + if port is not None: + self._servers.adb_port = port + return self + + def with_servers(self, servers: ServerConfig) -> "AgentConfigBuilder": + """Set the server settings. + + Args: + servers: The server settings to use + """ + self._servers = copy.deepcopy(servers) + return self + + def with_graph_config_callbacks(self, callbacks: Callbacks) -> "AgentConfigBuilder": + """Set the graph config callbacks. + + Args: + callbacks: The graph config callbacks to use + """ + self._graph_config_callbacks = callbacks + return self + + def with_video_recording_tools(self, enabled: bool = True) -> "AgentConfigBuilder": + """Enable or disable video recording tools. + + Args: + enabled: Whether to enable video recording tools + """ + self._video_recording_tools_enabled = enabled + return self + + def with_web_accessibility(self, enabled: bool = True) -> "AgentConfigBuilder": + """Enable or disable forcing web accessibility for WebViews. + + Args: + enabled: Whether to force web accessibility + """ + self._force_web_accessibility = enabled + return self + + def with_disable_checker(self, disable: bool = True) -> "AgentConfigBuilder": + """Temporarily disable the background Checker task (useful for debugging). + + Args: + disable: Whether to disable the checker + """ + self._disable_checker = disable + return self + + def with_checker( + self, + enabled: bool = True, + max_iterations: int | None = None, + max_chat_rounds: int | None = None, + ) -> "AgentConfigBuilder": + """Configure Checker subgoal verification and rollback settings. + + Args: + enabled: Whether checker visual verification is enabled + max_iterations: Optional maximum iterations for Checker reasoning (1-50) + max_chat_rounds: Optional maximum debate rounds with Operator (1-10) + """ + self._disable_checker = not enabled + if max_iterations is not None: + self._checker_max_iterations = max_iterations + if max_chat_rounds is not None: + self._checker_max_chat_rounds = max_chat_rounds + return self + + def with_disable_planner_validation(self, disable: bool = True) -> "AgentConfigBuilder": + """Temporarily disable the background Planner validation task. + + Args: + disable: Whether to disable the planner validation + """ + self._disable_planner_validation = disable + return self + + def with_planner_validation( + self, enabled: bool = True, similarity_threshold: float | None = None + ) -> "AgentConfigBuilder": + """Configure async Planner validation on milestone changes. + + Args: + enabled: Whether planner validation is enabled + similarity_threshold: Optional similarity threshold (0.0 - 1.0) + """ + self._disable_planner_validation = not enabled + if similarity_threshold is not None: + self._planner_validation_threshold = similarity_threshold + return self + + def with_enable_committee(self, enabled: bool = True) -> "AgentConfigBuilder": + """Enable or disable the Committee council debate tool for Operator.""" + self._enable_committee = enabled + return self + + def with_committee(self, enabled: bool = True, debate_rounds: int = 2) -> "AgentConfigBuilder": + """Configure Multi-Agent Committee council debate settings. + + Args: + enabled: Whether committee debate tool is enabled + debate_rounds: Number of debate rounds (1-5) + """ + self._enable_committee = enabled + self._committee_debate_rounds = debate_rounds + return self + + def with_outputter( + self, + enabled: bool = True, + force_synthesis: bool = False, + ) -> "AgentConfigBuilder": + """Configure Outputter post-execution synthesis agent options. + + Args: + enabled: Whether Outputter is mounted to synthesize final outputs + force_synthesis: Force Outputter synthesis even if no structured output schema is specified + """ + self._outputter = self._outputter.model_copy( + update={"enabled": enabled, "force_synthesis": force_synthesis} + ) + self._disable_outputter = not enabled + return self + + def with_disable_outputter(self, disable: bool = True) -> "AgentConfigBuilder": + """Disable Outputter agent execution.""" + self._disable_outputter = disable + self._outputter = self._outputter.model_copy(update={"enabled": not disable}) + return self + + def with_flash_config( + self, + max_turns: int | None = None, + explorer_mode: ExplorerVersion | None = None, + step_summarizer: bool | None = None, + step_summarizer_model: str | None = None, + prune_history_xml: bool | None = None, + ) -> "AgentConfigBuilder": + """Configure โšก Flash execution profile options. + + Args: + max_turns: Maximum reactive loop turns before timeout + explorer_mode: Visual perception mode for Flash runner ('flash' 1-shot detection) + step_summarizer: Enable/disable asynchronous visual context compressor + step_summarizer_model: Lightweight model for background step summarization + prune_history_xml: Whether to prune outdated XML trees from historical steps + """ + updates: dict[str, Any] = {} + if max_turns is not None: + updates["max_turns"] = max_turns + if explorer_mode is not None: + updates["explorer_mode"] = explorer_mode + self._explorer = self._explorer.model_copy(update={"flash_mode": explorer_mode}) + if ( + step_summarizer is not None + or step_summarizer_model is not None + or prune_history_xml is not None + ): + sum_updates: dict[str, Any] = {} + if step_summarizer is not None: + sum_updates["enabled"] = step_summarizer + if step_summarizer_model is not None: + sum_updates["model"] = step_summarizer_model + if prune_history_xml is not None: + sum_updates["prune_history_xml"] = prune_history_xml + updates["step_summarizer"] = self._flash.step_summarizer.model_copy(update=sum_updates) + if updates: + self._flash = self._flash.model_copy(update=updates) + return self + + def with_flash_step_summarizer( + self, + enabled: bool = True, + model: str | None = None, + prune_history_xml: bool | None = None, + ) -> "AgentConfigBuilder": + """Configure Flash asynchronous step state summarizer options.""" + return self.with_flash_config( + step_summarizer=enabled, + step_summarizer_model=model, + prune_history_xml=prune_history_xml, + ) + + def with_pro_config( + self, + explorer_mode: ExplorerVersion | None = None, + planner_validation: bool | None = None, + committee: bool | None = None, + checker: bool | None = None, + video_ledger: bool | None = None, + ) -> "AgentConfigBuilder": + """Configure ๐Ÿš€ Pro execution profile options. + + Args: + explorer_mode: Visual perception mode for Pro profile ('flash', 'pro', or 'ultra') + planner_validation: Enable/disable planner milestone validation + committee: Enable/disable multi-agent committee debate tool + checker: Enable/disable visual verification & snapshot rollback + video_ledger: Enable/disable screen video action ledger tracking + """ + if explorer_mode is not None: + self._pro = self._pro.model_copy( + update={"explorer": self._pro.explorer.model_copy(update={"mode": explorer_mode})} + ) + self._explorer = self._explorer.model_copy(update={"pro_mode": explorer_mode}) + if planner_validation is not None: + self.with_planner_validation(enabled=planner_validation) + if committee is not None: + self.with_committee(enabled=committee) + if checker is not None: + self.with_checker(enabled=checker) + if video_ledger is not None: + self._enable_video_ledger = video_ledger + self._pro = self._pro.model_copy( + update={ + "video_analyzer": self._pro.video_analyzer.model_copy( + update={"enable_ledger": video_ledger} + ) + } + ) + return self + + def with_explorer( + self, + version: ExplorerVersion | None = None, + default_version: ExplorerVersion | None = None, + flash_mode: ExplorerVersion | None = None, + pro_mode: ExplorerVersion | None = None, + caching: bool | None = None, + versions: dict[str, ExplorerVersion] | None = None, + ) -> "AgentConfigBuilder": + """Configure UI Explorer visual perception options. + + Args: + version: Default Explorer version mode ('flash', 'pro', or 'ultra') + default_version: Alias for version ('flash', 'pro', or 'ultra') + flash_mode: Explorer version for Flash execution profile (FlashRunner) + pro_mode: Explorer version for Pro execution profile (Graph/Operator/Validator) + caching: Whether to enable context caching for multi-turn Explorer + versions: Optional explicit per-agent version mapping + """ + target_version = version if version is not None else default_version + updates = {} + if target_version is not None: + updates["default_version"] = target_version + if flash_mode is not None: + updates["flash_mode"] = flash_mode + if pro_mode is not None: + updates["pro_mode"] = pro_mode + if caching is not None: + updates["caching"] = caching + if updates: + self._explorer = self._explorer.model_copy(update=updates) + if versions is not None: + self._explorer_versions = versions + return self + + def with_explorer_version(self, version: ExplorerVersion) -> "AgentConfigBuilder": + """Configure default model version to use for explorer sub-agents ('flash', 'pro', or 'ultra').""" + self._explorer = self._explorer.model_copy(update={"default_version": version}) + return self + + def with_explorer_versions(self, versions: dict[str, ExplorerVersion]) -> "AgentConfigBuilder": + """Configure per-agent model versions to use for explorer sub-agents.""" + self._explorer_versions = versions + return self + + def with_blacklisted_tools(self, tools: dict[str, list[str]]) -> "AgentConfigBuilder": + """Configure blacklisted tools.""" + self._blacklisted_tools = tools + return self + + def build(self, validate_profiles: bool = True) -> AgentConfig: + """Build the ARTEMIS AgentConfig object. + + Args: + default_profile: Name of the default agent profile to use + + Returns: + A configured AgentConfig object + + Raises: + ValueError: If default_profile is specified but not found in + configured profiles + """ + nb_profiles = len(self._agent_profiles) + + if isinstance(self._default_profile, str): + profile_name = self._default_profile + default_profile = self._agent_profiles.get(profile_name, None) + if default_profile is None: + raise ValueError(f"Profile '{profile_name}' not found in configured agents") + elif isinstance(self._default_profile, AgentProfile): + default_profile = self._default_profile + if default_profile.name not in self._agent_profiles: + self.add_profile(default_profile, validate=validate_profiles) + elif nb_profiles <= 0: + llm_config = get_default_llm_config() + default_profile = AgentProfile( + name=DEFAULT_PROFILE_NAME, + llm_config=llm_config, + ) + self.add_profile(default_profile, validate=validate_profiles) + elif nb_profiles == 1: + # Select the only one available + default_profile = next(iter(self._agent_profiles.values())) + else: + available_profiles = ", ".join(self._agent_profiles.keys()) + raise ValueError( + f"You must call with_default_profile() to select one among: {available_profiles}" + ) + + device_id = ( + self._device_id + or os.environ.get("ARTEMIS_DEVICE_ID") + or os.environ.get("ADB_DEVICE_SERIAL") + ) + + return AgentConfig( + agent_profiles=self._agent_profiles, + task_request_defaults=self._task_request_defaults or TaskRequestCommon(), + default_profile=default_profile, + device_id=device_id, + device_platform=self._device_platform, + servers=self._servers, + graph_config_callbacks=self._graph_config_callbacks, + video_recording_tools_enabled=self._video_recording_tools_enabled, + force_web_accessibility=self._force_web_accessibility, + disable_checker=self._disable_checker, + checker_max_iterations=self._checker_max_iterations, + checker_max_chat_rounds=self._checker_max_chat_rounds, + disable_planner_validation=self._disable_planner_validation, + planner_validation_threshold=self._planner_validation_threshold, + enable_committee=self._enable_committee, + committee_debate_rounds=self._committee_debate_rounds, + disable_outputter=self._disable_outputter, + outputter=self._outputter, + flash=self._flash, + pro=self._pro, + explorer=self._explorer, + explorer_versions=self._explorer_versions, + blacklisted_tools=self._blacklisted_tools, + enable_video_ledger=self._enable_video_ledger, + ) + + +def get_default_agent_config(): + return AgentConfigBuilder().build() + + +def get_default_servers(): + + host = settings.ADB_HOST or os.environ.get("ADB_HOST", "localhost") + port_val = settings.ADB_PORT + if not port_val: + port_str = os.environ.get("ADB_PORT", "5037") + port_val = int(port_str) if port_str.isdigit() else 5037 + + return ServerConfig( + adb_host=host, + adb_port=port_val, + ) diff --git a/artemis/sdk/builders/index.py b/artemis/sdk/builders/index.py new file mode 100644 index 0000000..7ecc3d1 --- /dev/null +++ b/artemis/sdk/builders/index.py @@ -0,0 +1,29 @@ +# 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. + +from artemis.sdk.builders.agent_config_builder import AgentConfigBuilder +from artemis.sdk.builders.task_request_builder import TaskRequestCommonBuilder + + +class BuildersWrapper: + @property + def AgentConfig(self) -> AgentConfigBuilder: + return AgentConfigBuilder() + + @property + def TaskDefaults(self) -> TaskRequestCommonBuilder: + return TaskRequestCommonBuilder() + + +Builders = BuildersWrapper() diff --git a/artemis/sdk/builders/task_request_builder.py b/artemis/sdk/builders/task_request_builder.py new file mode 100644 index 0000000..5f6e1e3 --- /dev/null +++ b/artemis/sdk/builders/task_request_builder.py @@ -0,0 +1,256 @@ +# 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. + +"""Builder for TaskRequest objects using a fluent interface.""" + +from pathlib import Path +from typing import Generic, TypeVar, cast + +try: + from typing import Self +except ImportError: + from typing import Self + +from artemis.constants import RECURSION_LIMIT +from artemis.sdk.types.agent import AgentProfile +from artemis.sdk.types.task import TaskRequest, TaskRequestCommon +from pydantic import BaseModel + +TIn = TypeVar("TIn", bound=BaseModel | None) +TOut = TypeVar("TOut", bound=BaseModel) + + +class _CyFunctionDetectorMeta(type): + def __instancecheck__(self, instance): + name = type(instance).__name__ + return ( + name + in ( + "cyfunction", + "cython_function_or_method", + "builtin_function_or_method", + ) + or "cyfunction" in name.lower() + ) + + +class CyFunctionDetector(metaclass=_CyFunctionDetectorMeta): + pass + + +class TaskRequestCommonBuilder(BaseModel): + """Builder class providing a fluent interface for creating TaskRequestCommon objects.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + + def __init__(self): + self._max_steps = RECURSION_LIMIT + self._record_trace = True + self._trace_path = Path("artemis-traces") + self._llm_output_path: Path | None = None + self._locked_app_package: str | None = None + self._app_path: Path | None = None + + def with_max_steps(self, max_steps: int) -> Self: + """Set the maximum number of steps the task can take. + + Args: + max_steps: Maximum number of steps + """ + self._max_steps = max_steps + return self + + def with_trace_recording(self, enabled: bool = True, path: str | None = None) -> Self: + """Configure trace recording for the task. + + Traces record screenshots and actions during execution. + + Args: + enabled: Whether to enable trace recording + path: Directory path where traces should be saved + """ + self._record_trace = enabled + if enabled and path: + self._trace_path = Path(path) + return self + + def with_llm_output_saving(self, path: str) -> Self: + """Configure LLM output saving for the task. + + Args: + path: Path where to save the LLM output message + """ + self._llm_output_path = Path(path) + return self + + def with_locked_app_package(self, package_name: str) -> Self: + """Set the app package to lock execution to. + + This ensures the specified app is launched and in the foreground before + the agentic loop starts. + + Args: + package_name: Package name (Android, e.g., 'com.whatsapp') + """ + self._locked_app_package = package_name + return self + + def with_app_path(self, app_path: str | Path) -> Self: + """Set the path to an app to install before running the task. + + For Android: Path to an APK file. + + The app will be installed automatically before the task starts. + + Args: + app_path: Path to the app file to install + """ + self._app_path = Path(app_path) if isinstance(app_path, str) else app_path + return self + + def build(self) -> TaskRequestCommon: + """Build the TaskRequestCommon object. + + Returns: + A configured TaskRequestCommon object + + Raises: + ValueError: If required fields are missing + """ + return TaskRequestCommon( + max_steps=self._max_steps, + record_trace=self._record_trace, + trace_path=self._trace_path, + llm_output_path=self._llm_output_path, + locked_app_package=self._locked_app_package, + app_path=self._app_path, + ) + + +class TaskRequestBuilder(TaskRequestCommonBuilder, Generic[TIn]): + """Builder class providing a fluent interface for creating TaskRequest objects. + + This builder allows for step-by-step construction of a TaskRequest with + clear methods that make the configuration process intuitive and type-safe. + + Examples: + >>> builder = TaskRequestBuilder[None](goal="Open Gmail and check unread + emails") + >>> task_request = ( + ... builder + ... .with_max_steps(30) + ... .using_profile("LowReasoning") + ... .with_output_description("A list of email subjects and senders") + ... .build() + ... ) + """ + + model_config = {"ignored_types": (CyFunctionDetector,)} + + def __init__(self, goal: str): + """Initialize an empty TaskRequestBuilder.""" + super().__init__() + self._goal = goal + self._profile: str | AgentProfile | None = None + self._name: str | None = None + self._output_description = None + self._output_format: type[TIn] | None = None + + @classmethod + def from_common(cls, goal: str, common: TaskRequestCommon): + res = cls(goal=goal) + res._max_steps = common.max_steps + res._record_trace = common.record_trace + res._trace_path = common.trace_path + res._llm_output_path = common.llm_output_path + res._locked_app_package = common.locked_app_package + res._app_path = common.app_path + return res + + def using_profile(self, profile: str | AgentProfile) -> "TaskRequestBuilder[TIn]": + """Set the agent profile for executing the task. + + Args: + profile: The agent profile to use + """ + self._profile = profile + return self + + def with_name(self, name: str) -> "TaskRequestBuilder[TIn]": + """Set the name of the task - useful when recording traces. + + Otherwise, a random name will be generated. + + Args: + name: Name of the task + """ + self._name = name + return self + + def without_llm_output_saving(self) -> Self: + """Disable LLM output saving for the task.""" + self._llm_output_path = None + return self + + def with_output_description(self, description: str) -> "TaskRequestBuilder[TIn]": + """Set the description of the expected output format. + + This is especially useful for data extraction tasks. + + Args: + description: Description of the expected output format + """ + self._output_description = description + return self + + def with_output_format(self, output_format: type[TOut]) -> "TaskRequestBuilder[TOut]": + """Set the pydantic model for the expected output format. + + Args: + output_format: Pydantic model instance defining the output format + """ + self._output_format = output_format # type: ignore + return cast(TaskRequestBuilder[TOut], self) + + def build(self) -> TaskRequest[TIn]: + """Build the TaskRequest object. + + Returns: + A configured TaskRequest object + + Raises: + ValueError: If required fields are missing + """ + if not self._goal: + raise ValueError("Task goal is required") + + if self._output_format and self._output_description: + raise ValueError("Output format and description are mutually exclusive") + + task_request = TaskRequest( + goal=self._goal, + profile=self._profile.name + if isinstance(self._profile, AgentProfile) + else self._profile, + task_name=self._name, + output_description=self._output_description, + output_format=self._output_format, + max_steps=self._max_steps, + record_trace=self._record_trace, + trace_path=self._trace_path, + llm_output_path=self._llm_output_path, + locked_app_package=self._locked_app_package, + app_path=self._app_path, + ) + return task_request diff --git a/artemis/sdk/constants.py b/artemis/sdk/constants.py new file mode 100644 index 0000000..99b32c5 --- /dev/null +++ b/artemis/sdk/constants.py @@ -0,0 +1,15 @@ +# 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. + +DEFAULT_PROFILE_NAME = "default" diff --git a/artemis/sdk/examples/__init__.py b/artemis/sdk/examples/__init__.py new file mode 100644 index 0000000..d8e7628 --- /dev/null +++ b/artemis/sdk/examples/__init__.py @@ -0,0 +1,15 @@ +# 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. + +"""Example scripts for the ARTEMIS SDK.""" diff --git a/artemis/sdk/examples/app_lock_messaging.py b/artemis/sdk/examples/app_lock_messaging.py new file mode 100644 index 0000000..e89e7fb --- /dev/null +++ b/artemis/sdk/examples/app_lock_messaging.py @@ -0,0 +1,87 @@ +# 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 asyncio + +from artemis.sdk import Agent +from pydantic import BaseModel, Field + + +class _CyFunctionDetectorMeta(type): + def __instancecheck__(self, instance): + name = type(instance).__name__ + return ( + name + in ( + "cyfunction", + "cython_function_or_method", + "builtin_function_or_method", + ) + or "cyfunction" in name.lower() + ) + + +class CyFunctionDetector(metaclass=_CyFunctionDetectorMeta): + pass + + +class MessageResult(BaseModel): + """Structured result from messaging task.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + + messages_sent: int = Field(..., description="Number of messages successfully sent") + contacts: list[str] = Field(..., description="List of contacts messaged") + success: bool = Field(..., description="Whether all messages were sent successfully") + + +async def main() -> None: + # Create agent with default configuration + agent = Agent() + + try: + await agent.init() + + # Use app lock to keep execution in WhatsApp + # This ensures the agent stays in the app and relaunches if needed + task = ( + agent.new_task("Send 'Happy New Year!' message to Alice, Bob, and Charlie on WhatsApp") + .with_name("send_new_year_messages") + .with_locked_app_package("com.whatsapp") # Lock to WhatsApp + .with_output_format(MessageResult) + .with_max_steps(600) # Messaging tasks may need more steps + .build() + ) + + print("Sending messages with app lock enabled...") + print("The agent will stay in WhatsApp and relaunch if needed.\n") + + result = await agent.run_task(request=task) + + if result: + print("\n=== Messaging Complete ===") + print(f"Messages sent: {result.messages_sent}") + print(f"Contacts: {', '.join(result.contacts)}") + print(f"Success: {result.success}") + else: + print("Failed to send messages") + + except Exception as e: + print(f"Error: {e}") + finally: + await agent.clean() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/artemis/sdk/examples/simple_photo_organizer.py b/artemis/sdk/examples/simple_photo_organizer.py new file mode 100644 index 0000000..6650d66 --- /dev/null +++ b/artemis/sdk/examples/simple_photo_organizer.py @@ -0,0 +1,111 @@ +# 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. + +"""Simple Photo Organizer - Basic SDK Usage Example + +This example demonstrates a straightforward way to use the Artemis SDK +without builders or advanced configuration. It performs a real-world automation +task: +1. Opens the photo gallery +2. Finds photos from a specific date +3. Creates an album and moves those photos into it + +Run: +- python artemis/sdk/examples/simple_photo_organizer.py +""" + +import asyncio +from datetime import date, timedelta +from artemis.sdk import Agent +from pydantic import BaseModel, Field + + +class _CyFunctionDetectorMeta(type): + def __instancecheck__(self, instance): + name = type(instance).__name__ + return ( + name + in ( + "cyfunction", + "cython_function_or_method", + "builtin_function_or_method", + ) + or "cyfunction" in name.lower() + ) + + +class CyFunctionDetector(metaclass=_CyFunctionDetectorMeta): + pass + + +class PhotosResult(BaseModel): + """Structured result from photo search.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + + found_photos: int = Field(..., description="Number of photos found") + date_range: str = Field(..., description="Date range of photos found") + album_created: bool = Field(..., description="Whether an album was created") + album_name: str = Field(..., description="Name of the created album") + photos_moved: int = Field(0, description="Number of photos moved to the album") + + +async def main() -> None: + # Create a simple agent with default configuration + agent = Agent() + + try: + # Initialize agent (finds a device, starts required servers) + await agent.init() + + # Calculate yesterday's date for the example + yesterday = date.today() - timedelta(days=1) + formatted_date = yesterday.strftime("%B %d") # e.g. "August 22" + + print(f"Looking for photos from {formatted_date}...") + + # First task: search for photos and organize them, with typed output + result = await agent.run_task( + goal=( + "Open the Photos/Gallery app. Find photos taken on" + f" {formatted_date}. Create a new album named '{formatted_date}" + " Memories' and move those photos into it. Count how many" + " photos were moved." + ), + output=PhotosResult, + name="organize_photos", + ) + + # Handle and display the result + if result: + print("\n=== Photo Organization Complete ===") + print(f"Found: {result.found_photos} photos from {result.date_range}") + + if result.album_created: + print(f"Created album: '{result.album_name}'") + print(f"Moved {result.photos_moved} photos to the album") + else: + print("No album was created") + else: + print("Failed to organize photos") + + except Exception as e: + print(f"Error: {e}") + finally: + # Always clean up resources + await agent.clean() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/artemis/sdk/examples/smart_notification_assistant.py b/artemis/sdk/examples/smart_notification_assistant.py new file mode 100644 index 0000000..297ea08 --- /dev/null +++ b/artemis/sdk/examples/smart_notification_assistant.py @@ -0,0 +1,235 @@ +# 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. + +"""Smart Notification Assistant - Intermediate SDK Usage Example + +This example demonstrates more advanced SDK features including: +- TaskRequestBuilder pattern +- Multiple agent profiles for different reasoning tasks +- Tracing for debugging/visualization +- Structured output with Pydantic +- Exception handling + +It performs a practical automation task: +1. Checks notification panel for unread notifications +2. Categorizes them by priority/app +3. Performs actions based on notification content + +Run: +- python artemis/sdk/examples/smart_notification_assistant.py +""" + +import asyncio +from datetime import datetime + +try: + from enum import StrEnum +except ImportError: + from enum import Enum + + class StrEnum(str, Enum): + pass + + +from artemis.config import LLM, LLMConfig, LLMConfigUtils, LLMWithFallback +from artemis.sdk import Agent +from artemis.sdk.builders import Builders +from artemis.sdk.types import AgentProfile +from artemis.sdk.types.exceptions import AgentError +from pydantic import BaseModel, Field + + +class NotificationPriority(StrEnum): + HIGH = "high" + MEDIUM = "medium" + LOW = "low" + + +class _CyFunctionDetectorMeta(type): + def __instancecheck__(self, instance): + name = type(instance).__name__ + return ( + name + in ( + "cyfunction", + "cython_function_or_method", + "builtin_function_or_method", + ) + or "cyfunction" in name.lower() + ) + + +class CyFunctionDetector(metaclass=_CyFunctionDetectorMeta): + pass + + +class Notification(BaseModel): + """Individual notification details.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + + app_name: str = Field(..., description="Name of the app that sent the notification") + title: str = Field(..., description="Title/header of the notification") + message: str = Field(..., description="Message content of the notification") + priority: NotificationPriority = Field( + default=NotificationPriority.MEDIUM, + description="Priority level of notification", + ) + + +class NotificationSummary(BaseModel): + """Summary of all notifications.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + + total_count: int = Field(..., description="Total number of notifications found") + high_priority_count: int = Field(0, description="Count of high priority notifications") + notifications: list[Notification] = Field( + default_factory=list, description="List of individual notifications" + ) + + +def get_agent() -> Agent: + # Create two specialized profiles: + # 1. An analyzer profile for detailed inspection tasks + analyzer_profile = AgentProfile( + name="analyzer", + llm_config=LLMConfig( + planner=LLMWithFallback( + provider="openrouter", + model="meta-llama/llama-4-scout", + fallback=LLM(provider="openrouter", model="meta-llama/llama-4-maverick"), + ), + utils=LLMConfigUtils( + outputter=LLMWithFallback( + provider="openai", + model="gpt-5-nano", + fallback=LLM(provider="openai", model="gpt-5-mini"), + ), + hopper=LLMWithFallback( + provider="openai", + model="gpt-5-nano", + fallback=LLM(provider="openai", model="gpt-5-mini"), + ), + ), + ), + # from_file="/tmp/analyzer.jsonc" # can be loaded from file + ) + + # 2. An action profile for handling easy & fast actions based on notifications + action_profile = AgentProfile( + name="note_taker", + llm_config=LLMConfig( + planner=LLMWithFallback( + provider="openai", + model="o3", + fallback=LLM(provider="openai", model="gpt-5"), + ), + utils=LLMConfigUtils( + outputter=LLMWithFallback( + provider="openai", + model="gpt-5-nano", + fallback=LLM(provider="openai", model="gpt-5-mini"), + ), + hopper=LLMWithFallback( + provider="openai", + model="gpt-5-nano", + fallback=LLM(provider="openai", model="gpt-5-mini"), + ), + ), + ), + ) + + # Configure default task settings with tracing + task_defaults = Builders.TaskDefaults.with_max_steps(200).build() + + # Configure the agent + config = ( + Builders.AgentConfig.add_profiles(profiles=[analyzer_profile, action_profile]) + .with_default_profile(profile=action_profile) + .with_default_task_config(config=task_defaults) + .build() + ) + return Agent(config=config) + + +async def main(): + # Set up traces directory with timestamp for uniqueness + timestamp = datetime.now().strftime("%Y%m%d_%H%M") + traces_dir = f"/tmp/notification_traces/{timestamp}" + agent = get_agent() + + try: + # Initialize agent (finds a device, starts required servers) + await agent.init() + + print("Checking for notifications...") + + # Task 1: Get and analyze notifications with analyzer profile + notification_task = ( + agent.new_task( + goal=( + "Open the notification panel (swipe down from top). Scroll" + " through the first 3 unread notifications. For each" + " notification, identify the app name, title, and content." + " Tag messages from messaging apps or email as high" + " priority." + ) + ) + .with_output_format(NotificationSummary) + .using_profile("analyzer") + .with_name("notification_scan") + .with_max_steps(400) + .with_trace_recording(enabled=True, path=traces_dir) + .build() + ) + + # Execute the task with proper exception handling + try: + notifications = await agent.run_task(request=notification_task) + + # Display the structured results + if notifications: + print("\n=== Notification Summary ===") + print(f"Total notifications: {notifications.total_count}") + print(f"High priority: {notifications.high_priority_count}") + + # Task 2: Create a note to store the notification summary + response = await agent.run_task( + goal=( + "Open my Notes app and create a new note summarizing" + f" the following information:\n{notifications}" + ), + name="email_action", + profile="note_taker", + ) + print(f"Action result: {response}") + + else: + print("Failed to retrieve notifications") + + except AgentError as e: + print(f"Agent error occurred: {e}") + except Exception as e: + print(f"Unexpected error: {type(e).__name__}: {e}") + raise + + finally: + # Clean up + await agent.clean() + print(f"\nTraces saved to: {traces_dir}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/artemis/sdk/examples/video_transcription_example.py b/artemis/sdk/examples/video_transcription_example.py new file mode 100644 index 0000000..d69f1b8 --- /dev/null +++ b/artemis/sdk/examples/video_transcription_example.py @@ -0,0 +1,109 @@ +# 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. + +"""Video Transcription Example + +This example demonstrates how to use the video recording tools to capture +and analyze video content from a mobile device screen. + +The agent can: +1. Start a screen recording +2. Perform actions while recording +3. Stop the recording and analyze its content using Gemini models + +Use case: Recording a video playing on the screen and transcribing its content. +""" + +import asyncio + +from artemis.config import LLM, LLMConfig, LLMConfigUtils, LLMWithFallback +from artemis.sdk.agent import Agent +from artemis.sdk.builders.agent_config_builder import AgentConfigBuilder +from artemis.sdk.types.agent import AgentConfig +from artemis.sdk.types.task import AgentProfile, TaskRequest + + +def get_video_capable_llm_config() -> LLMConfig: + """Returns an LLM config with video_analyzer configured. + + The video_analyzer must use a video-capable Gemini model: + - gemini-3.6-flash (recommended - fast and capable) + - gemini-3-pro-preview + - gemini-2.5-flash + - gemini-2.5-pro + - gemini-2.0-flash + """ + return LLMConfig( + planner=LLMWithFallback( + provider="openai", + model="gpt-5-nano", + fallback=LLM(provider="openai", model="gpt-5-mini"), + ), + utils=LLMConfigUtils( + outputter=LLMWithFallback( + provider="openai", + model="gpt-5-nano", + fallback=LLM(provider="openai", model="gpt-5-mini"), + ), + hopper=LLMWithFallback( + provider="openai", + model="gpt-5-nano", + fallback=LLM(provider="openai", model="gpt-5-mini"), + ), + video_analyzer=LLMWithFallback( + provider="google", + model="gemini-3.6-flash", + fallback=LLM(provider="google", model="gemini-2.5-flash"), + ), + ), + ) + + +async def main(): + config: AgentConfig = ( + AgentConfigBuilder() + .add_profile( + AgentProfile( + name="VideoCapable", + llm_config=get_video_capable_llm_config(), + ) + ) + .with_video_recording_tools() + .build() + ) + + agent = Agent(config=config) + try: + await agent.init() + + result = await agent.run_task( + request=TaskRequest( + goal=""" + 1. Open YouTube app + 2. Search for "Python tutorial" + 3. Start recording the screen + 4. Play the first video + 5. Wait for the first 30 seconds of the video to play + 6. Stop recording and tell me what was said in the video + """, + profile="VideoCapable", + ) + ) + print(f"Task result: {result}") + finally: + await agent.clean() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/artemis/sdk/types/__init__.py b/artemis/sdk/types/__init__.py new file mode 100644 index 0000000..9c59ed1 --- /dev/null +++ b/artemis/sdk/types/__init__.py @@ -0,0 +1,61 @@ +# 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. + +"""Type definitions for the ARTEMIS SDK.""" + +from artemis.sdk.types.agent import ( + AgentConfig, + ApiBaseUrl, + DevicePlatform, + ServerConfig, +) +from artemis.sdk.types.exceptions import ( + AgentError, + AgentNotInitializedError, + AgentProfileNotFoundError, + AgentTaskRequestError, + DeviceError, + DeviceNotFoundError, + ArtemisError, + ServerError, + ServerStartupError, +) +from artemis.sdk.types.task import ( + AgentProfile, + Task, + TaskRequest, + TaskRequestCommon, + TaskResult, +) + +__all__ = [ + "ApiBaseUrl", + "AgentConfig", + "DevicePlatform", + "AgentProfile", + "ServerConfig", + "TaskRequest", + "TaskResult", + "TaskRequestCommon", + "Task", + "AgentProfileNotFoundError", + "AgentTaskRequestError", + "DeviceNotFoundError", + "ServerStartupError", + "AgentError", + "AgentNotInitializedError", + "DeviceError", + "ArtemisError", + "ServerError", +] diff --git a/artemis/sdk/types/agent.py b/artemis/sdk/types/agent.py new file mode 100644 index 0000000..090c8cf --- /dev/null +++ b/artemis/sdk/types/agent.py @@ -0,0 +1,145 @@ +# 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. + +from typing import Literal +from urllib.parse import urlparse + +from langchain_core.callbacks.base import Callbacks +from artemis.config import ( + ExplorerConfig, + FlashProfileConfig, + OutputterConfig, + ProProfileConfig, +) +from artemis.context import DevicePlatform +from artemis.sdk.types.task import AgentProfile, TaskRequestCommon +from artemis.utils.video import detect_video_tools_enabled +from pydantic import BaseModel, Field + + +class _CyFunctionDetectorMeta(type): + def __instancecheck__(self, instance): + name = type(instance).__name__ + return ( + name + in ( + "cyfunction", + "cython_function_or_method", + "builtin_function_or_method", + ) + or "cyfunction" in name.lower() + ) + + +class CyFunctionDetector(metaclass=_CyFunctionDetectorMeta): + pass + + +class ApiBaseUrl(BaseModel): + """Defines an API base URL.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + scheme: Literal["http", "https"] + host: str + port: int | None = None + + def __eq__(self, other): + if not isinstance(other, ApiBaseUrl): + return False + return self.to_url() == other.to_url() + + def to_url(self): + return ( + f"{self.scheme}://{self.host}:{self.port}" + if self.port is not None + else f"{self.scheme}://{self.host}" + ) + + @classmethod + def from_url(cls, url: str) -> "ApiBaseUrl": + parsed_url = urlparse(url) + if parsed_url.scheme not in ["http", "https"]: + raise ValueError(f"Invalid scheme: {parsed_url.scheme}") + if parsed_url.hostname is None: + raise ValueError("Invalid hostname") + return cls( + scheme=parsed_url.scheme, # type: ignore + host=parsed_url.hostname, + port=parsed_url.port, + ) + + +class ServerConfig(BaseModel): + """Configuration for the required servers.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + + adb_host: str + adb_port: int + + +class AgentConfig(BaseModel): + """ARTEMIS agent configuration. + + Attributes: + agent_profiles: Map an agent profile name to its configuration. + task_config_defaults: Default task request configuration. + default_profile: default profile to use for tasks + device_id: Specific device to target (if None, first available is used). + device_platform: Platform of the device to target. + servers: Custom server configurations. + """ + + agent_profiles: dict[str, AgentProfile] + task_request_defaults: TaskRequestCommon + default_profile: AgentProfile + device_id: str | None = None + device_platform: DevicePlatform | None = None + servers: ServerConfig + graph_config_callbacks: Callbacks = None + video_recording_tools_enabled: bool = Field(default_factory=detect_video_tools_enabled) + force_web_accessibility: bool = False + disable_checker: bool = True + checker_max_iterations: int = 20 + checker_max_chat_rounds: int = 4 + disable_planner_validation: bool = True + planner_validation_threshold: float = 0.85 + enable_committee: bool = False + committee_debate_rounds: int = 2 + disable_outputter: bool = False + outputter: OutputterConfig = Field(default_factory=OutputterConfig) + + flash: FlashProfileConfig = Field(default_factory=FlashProfileConfig) + pro: ProProfileConfig = Field(default_factory=ProProfileConfig) + + explorer: ExplorerConfig = Field(default_factory=ExplorerConfig) + explorer_versions: dict[str, Literal["flash", "pro", "ultra"]] + blacklisted_tools: dict[str, list[str]] + enable_video_ledger: bool = True + + model_config = {"arbitrary_types_allowed": True} + + def get_explorer_version( + self, + explicit_version: str | None = None, + agent_name: str | None = "operator", + profile: str | None = None, + ) -> Literal["flash", "pro", "ultra"]: + """Resolves active Explorer version using underlying ExplorerConfig and role overrides.""" + return self.explorer.resolve( + explicit_version=explicit_version, + agent_name=agent_name, + profile=profile, + per_agent_overrides=self.explorer_versions, + ) diff --git a/artemis/sdk/types/exceptions.py b/artemis/sdk/types/exceptions.py new file mode 100644 index 0000000..63288cc --- /dev/null +++ b/artemis/sdk/types/exceptions.py @@ -0,0 +1,105 @@ +# 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. + +"""Exceptions for the ARTEMIS SDK. + +This module defines the exception hierarchy used throughout the Artemis SDK. +""" + +from typing import Literal + + +class ArtemisError(Exception): + """Base exception class for all ARTEMIS SDK exceptions.""" + + def __init__(self, message="An error occurred in the Artemis SDK"): + self.message = message + super().__init__(self.message) + + +class DeviceError(ArtemisError): + """Exception raised for errors related to mobile devices.""" + + def __init__(self, message="A device-related error occurred"): + super().__init__(message) + + +class DeviceNotFoundError(DeviceError): + """Exception raised when no mobile device is found.""" + + def __init__(self, message="No mobile device found"): + super().__init__(message) + + +class ServerError(ArtemisError): + """Exception raised for errors related to ARTEMIS servers.""" + + def __init__(self, message="A server-related error occurred"): + super().__init__(message) + + +class ServerStartupError(ServerError): + """Exception raised when ARTEMIS servers fail to start.""" + + def __init__(self, server_name=None, message=None): + if server_name and not message: + message = f"Failed to start {server_name}" + elif not message: + message = "Failed to start Artemis servers" + super().__init__(message) + self.server_name = server_name + + +class AgentError(ArtemisError): + """Exception raised for errors related to the ARTEMIS agent.""" + + def __init__(self, message="An agent-related error occurred"): + super().__init__(message) + + +class AgentNotInitializedError(AgentError): + """Exception raised when attempting operations on an uninitialized agent.""" + + def __init__(self, message="Agent is not initialized. Call init() first"): + super().__init__(message) + + +class AgentTaskRequestError(AgentError): + """Exception raised when a requested task is invalid.""" + + def __init__(self, message="An agent task-related error occurred"): + super().__init__(message) + + +class AgentProfileNotFoundError(AgentTaskRequestError): + """Exception raised when an agent profile is not found.""" + + def __init__(self, profile_name: str): + super().__init__(f"Agent profile {profile_name} not found") + + +EXECUTABLES = Literal["adb"] + + +class ExecutableNotFoundError(ArtemisError): + """Exception raised when a required executable is not found.""" + + def __init__(self, executable_name: EXECUTABLES): + install_instructions: dict[EXECUTABLES, str] = { + "adb": "https://developer.android.com/tools/adb", + } + message = f"Required executable '{executable_name}' not found in PATH." + if executable_name in install_instructions: + message += f"\nTo install it, please visit: {install_instructions[executable_name]}" + super().__init__(message) diff --git a/artemis/sdk/types/platform.py b/artemis/sdk/types/platform.py new file mode 100644 index 0000000..5bceb4f --- /dev/null +++ b/artemis/sdk/types/platform.py @@ -0,0 +1,210 @@ +# 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. + +from datetime import datetime +from typing import Annotated, Any, Literal + +from pydantic import BaseModel, ConfigDict, Field, PlainSerializer +from pydantic.v1.utils import to_lower_camel + +TaskRunStatus = Literal["pending", "running", "completed", "failed", "cancelled"] + +IsoDatetime = Annotated[ + datetime, + PlainSerializer( + func=lambda v: v.isoformat() if v else None, + return_type=str, + when_used="unless-none", + ), +] + + +class _CyFunctionDetectorMeta(type): + def __instancecheck__(self, instance): + name = type(instance).__name__ + return ( + name + in ( + "cyfunction", + "cython_function_or_method", + "builtin_function_or_method", + ) + or "cyfunction" in name.lower() + ) + + +class CyFunctionDetector(metaclass=_CyFunctionDetectorMeta): + pass + + +class BaseApiModel(BaseModel): + model_config = ConfigDict( + alias_generator=to_lower_camel, + populate_by_name=True, + str_strip_whitespace=True, + ) + + +class LLMProfileResponse(BaseApiModel): + """Response model for LLM profile.""" + + id: str = Field(..., description="Profile ID") + name: str = Field(..., description="Profile name") + description: str | None = Field(None, description="Profile description") + llms: dict[str, Any] = Field(..., description="LLM configuration") + created_at: str = Field(..., description="Creation timestamp (ISO format)") + updated_at: str = Field(..., description="Last update timestamp (ISO format)") + + +class TaskOptionsResponse(BaseApiModel): + """Response model for task options.""" + + id: str = Field(..., description="Options ID") + enable_tracing: bool = Field(..., description="Whether tracing is enabled") + max_steps: int = Field(..., description="Maximum number of steps") + enable_video_tools: bool = Field(..., description="Whether video tools are enabled") + created_at: str = Field(..., description="Creation timestamp (ISO format)") + updated_at: str = Field(..., description="Last update timestamp (ISO format)") + + +class TaskResponse(BaseApiModel): + """Response model for task.""" + + id: str = Field(..., description="Task ID") + name: str = Field(..., description="Task name") + description: str | None = Field(None, description="Task description") + input_prompt: str = Field(..., description="Input prompt") + output_description: str | None = Field(None, description="Output description") + options: TaskOptionsResponse = Field(..., description="Task options") + created_at: str = Field(..., description="Creation timestamp (ISO format)") + updated_at: str = Field(..., description="Last update timestamp (ISO format)") + + +class CreateTaskRunRequest(BaseApiModel): + """Request model for creating a task run.""" + + task_id: str = Field(..., description="ID of the task to run") + llm_profile_id: str = Field(..., description="LLM profile ID to use") + virtual_mobile_id: str | None = Field(None, description="Virtual mobile ID to use") + locked_app_package: str | None = Field(None, description="App package to lock for the task run") + execution_origin: str | None = Field(None, description="Origin of the task execution") + enable_video_tools: bool = Field( + default=False, + description="Whether video tools are enabled for this task run", + ) + + +class CreateOrphanTaskRunRequest(BaseApiModel): + """Request model for creating an orphan task run (without a parent task).""" + + input_prompt: str = Field( + ..., + description="Input prompt for the task run", + ) + output_description: str | None = Field( + None, + description="Optional output description", + ) + llm_profile_id: str = Field( + ..., + description="LLM profile ID to use", + ) + virtual_mobile_id: str | None = Field( + None, + description="Optional virtual mobile ID on which it runs", + ) + locked_app_package: str | None = Field( + None, + description="Optional app package to lock execution to", + ) + execution_origin: str | None = Field( + None, + description="Origin of the task run execution", + ) + max_steps: int | None = Field( + None, + description="Maximum number of steps for the task run", + ) + enable_video_tools: bool = Field( + default=False, + description="Whether video tools are enabled for this task run", + ) + + +class UpdateTaskRunStatusRequest(BaseApiModel): + """Request model for updating task run status.""" + + status: TaskRunStatus = Field(..., description="New status of the task run") + message: str | None = Field(None, description="Message associated with the status") + output: str | None = Field(None, description="Output of the task run") + + +class TaskRunResponse(BaseApiModel): + """Response model for a single task run.""" + + id: str = Field(..., description="Unique identifier for the task run") + task: TaskResponse | None = Field( + ..., + description=("ID of the task this run is for or None if manually created"), + ) + llm_profile: LLMProfileResponse = Field(..., description="LLM profile ID used for this run") + status: TaskRunStatus = Field(..., description="Current status of the task run") + input_prompt: str = Field(..., description="Input prompt for this task run") + output_description: str | None = Field(None, description="Description of expected output") + created_at: datetime = Field(..., description="When the task run was created") + started_at: datetime | None = Field(None, description="When the task run started") + finished_at: datetime | None = Field(None, description="When the task run finished") + locked_app_package: str | None = Field( + None, description="The app package to lock for the task run" + ) + + +SubgoalState = Literal["pending", "started", "completed", "failed"] + + +class ArtemisSubgoal(BaseModel): + """Upsert ArtemisSubgoal API model.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + + name: str = Field(..., description="Name of the subgoal") + state: SubgoalState = Field(default="pending", description="Current state of the subgoal") + started_at: IsoDatetime | None = Field(default=None, description="When the subgoal started") + ended_at: IsoDatetime | None = Field(default=None, description="When the subgoal ended") + + +class UpsertTaskRunPlanRequest(BaseApiModel): + """Upsert task run plan request model.""" + + started_at: IsoDatetime = Field(..., description="When the plan started") + subgoals: list[ArtemisSubgoal] = Field(..., description="Subgoals of the plan") + ended_at: IsoDatetime | None = Field( + default=None, + description="When the plan ended (replanned or completed)", + ) + + +class TaskRunPlanResponse(UpsertTaskRunPlanRequest): + """Response model for a task run plan.""" + + id: str = Field(..., description="Unique identifier for the task run plan") + task_run_id: str = Field(..., description="ID of the task run this plan is for") + + +class TrajectoryGifUploadResponse(BaseApiModel): + """Response containing trajectory GIF signed upload URL.""" + + signed_url: str = Field(..., description="Signed upload URL for the trajectory GIF") + task_run_id: str = Field(..., description="Task run ID for which the GIF URL was generated") + file_path: str = Field(..., description="Path where the GIF will be stored") diff --git a/artemis/sdk/types/task.py b/artemis/sdk/types/task.py new file mode 100644 index 0000000..a57cf8e --- /dev/null +++ b/artemis/sdk/types/task.py @@ -0,0 +1,246 @@ +# 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. + +"""Task-related type definitions for the ARTEMIS SDK.""" + +from collections.abc import Callable, Coroutine +from datetime import datetime +from pathlib import Path +from typing import Any, Generic, TypeVar, overload + +from artemis.config import LLMConfig, get_default_llm_config +from artemis.constants import RECURSION_LIMIT +from artemis.context import DeviceContext +from artemis.sdk.types.platform import TaskRunStatus +from artemis.sdk.utils import load_llm_config_override +from pydantic import BaseModel, Field + + +class _CyFunctionDetectorMeta(type): + def __instancecheck__(self, instance): + name = type(instance).__name__ + return ( + name + in ( + "cyfunction", + "cython_function_or_method", + "builtin_function_or_method", + ) + or "cyfunction" in name.lower() + ) + + +class CyFunctionDetector(metaclass=_CyFunctionDetectorMeta): + pass + + +class AgentProfile(BaseModel): + """Represents a ARTEMIS agent profile. + + Attributes: + name: Name of the agent - used to reference the agent when running + tasks. + llm_config: LLM configuration for the agent. + """ + + model_config = {"ignored_types": (CyFunctionDetector,)} + + name: str + llm_config: LLMConfig = Field(default_factory=get_default_llm_config) + + @overload + def __init__(self, *, name: str, llm_config: LLMConfig): ... + + @overload + def __init__(self, *, name: str, from_file: str): ... + + def __init__( + self, + *, + name: str, + llm_config: LLMConfig | None = None, + from_file: str | None = None, + **kwargs, + ): + kwargs["name"] = name + if from_file: + kwargs["llm_config"] = load_llm_config_override(Path(from_file)) + elif llm_config: + kwargs["llm_config"] = llm_config + else: + raise ValueError("Either llm_config or from_file must be provided") + super().__init__(**kwargs) + + def __str__(self): + return f"Profile {self.name}:\n{self.llm_config}" + + +T = TypeVar("T", bound=BaseModel) +TOutput = TypeVar("TOutput", bound=BaseModel | None) + + +class TaskRequestBase(BaseModel): + """Defines common parameters of a mobile automation task request.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + max_steps: int = RECURSION_LIMIT + record_trace: bool = True + trace_path: Path = Path("artemis-traces") + llm_output_path: Path | None = None + + +class TaskRequestCommon(TaskRequestBase): + """Defines common parameters for any task request.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + max_steps: int = RECURSION_LIMIT + locked_app_package: str | None = None + app_path: Path | None = None + """Path to an app to install before running the task. + + For Android: Path to an APK file. + + The app will be installed automatically before the task starts. + """ + + +class TaskRequest(TaskRequestCommon, Generic[TOutput]): + """Defines the format of a mobile automation task request. + + Attributes: + goal: Natural language description of the goal to achieve + profile: Optional agent profile to use for executing the task + task_name: Optional name for the task + output_description: Optional natural language description of expected + output format + output_format: Optional pydantic model for the output format of the task + max_steps: Maximum number of steps the agent can take (default: 20) + record_trace: Whether to record a trace (screenshots, actions) of the + execution (default: False) + trace_path: Directory path to save trace data if recording is enabled + llm_output_path: Path to save LLM output data + """ + + model_config = {"ignored_types": (CyFunctionDetector,)} + goal: str + profile: str | None = None + task_name: str | None = None + output_description: str | None = None + output_format: type[TOutput] | None = None + enable_remote_tracing: bool = False + + +class TaskResult(BaseModel): + """Result of a mobile automation task. + + Attributes: + content: Raw result content (could be text or structured data) + error: Error message if the task failed + execution_time_seconds: How long the task took to execute + steps_taken: Number of steps the agent took to complete the task + """ + + model_config = {"ignored_types": (CyFunctionDetector,)} + content: Any = None + error: str | None = None + execution_time_seconds: float + steps_taken: int + + def get_as_model(self, model_class: type[T]) -> T: + """Parse the content into a Pydantic model instance. + + Args: + model_class: The Pydantic model class to parse the data into + + Returns: + An instance of the specified model class + + Raises: + ValueError: If content is None or not compatible with the model + """ + if self.content is None: + raise ValueError("No content available to parse into a model") + if isinstance(self.content, model_class): + return self.content + return model_class.model_validate(self.content) + + +class Task(BaseModel): + """A mobile automation task to be executed. + + Attributes: + id: Unique identifier for the task + device: Information about the target device + status: Current status of the task execution + request: User task request + created_at: ISO timestamp when the task was created + ended_at: ISO timestamp when the task ended + """ + + model_config = {"ignored_types": (CyFunctionDetector,)} + id: str + device: DeviceContext + status: TaskRunStatus + status_message: str | None = None + on_status_changed: Callable[[TaskRunStatus, str | None, Any | None], Coroutine] | None = None + request: TaskRequest + created_at: datetime + ended_at: datetime | None = None + result: TaskResult | None = None + + async def finalize( + self, + content: Any | None = None, + state: dict | None = None, + error: str | None = None, + cancelled: bool = False, + ): + new_status: TaskRunStatus = "completed" if error is None else "failed" + if new_status == "failed" and cancelled: + new_status = "cancelled" + message = "Task completed successfully" + if new_status == "failed": + message = "Task failed" + (f": {error}" if error else "") + elif new_status == "cancelled": + message = "Task cancelled" + (f": {error}" if error else "") + await self.set_status(status=new_status, message=message, output=content or error) + self.ended_at = datetime.now() + + duration = self.ended_at - self.created_at + steps_taken = -1 + if state is not None: + metadata = state.get("metadata", None) + if metadata: + steps_taken = metadata.get("step_count", -1) + + self.result = TaskResult( + content=content, + error=error, + execution_time_seconds=duration.total_seconds(), + steps_taken=steps_taken, + ) + + def get_name(self) -> str: + return self.request.task_name or self.id + + async def set_status( + self, + status: TaskRunStatus, + message: str | None = None, + output: Any | None = None, + ): + self.status = status + self.status_message = message + if self.on_status_changed: + await self.on_status_changed(status, message, output) diff --git a/artemis/sdk/utils.py b/artemis/sdk/utils.py new file mode 100644 index 0000000..43ed45b --- /dev/null +++ b/artemis/sdk/utils.py @@ -0,0 +1,28 @@ +# 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. + + +from artemis.config import ( + LLMConfig, + deep_merge_llm_config, + get_default_llm_config, + load_llm_config_override, +) + +__all__ = [ + "LLMConfig", + "deep_merge_llm_config", + "get_default_llm_config", + "load_llm_config_override", +] diff --git a/artemis/services/llm.py b/artemis/services/llm.py new file mode 100644 index 0000000..f1e1df4 --- /dev/null +++ b/artemis/services/llm.py @@ -0,0 +1,635 @@ +# 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. + +"""Universal LLM Service for Artemis. + +Provides connection-pooled model instantiation, robust retries, +thought stream recording, and role-based dynamic dispatching. +""" + +import asyncio +from collections.abc import Awaitable, Callable, Coroutine +from dataclasses import replace +import functools +import logging +from pathlib import Path +import random +from typing import Any, Literal, TypeVar, overload +from uuid import uuid4 + +from langchain_core.language_models.chat_models import BaseChatModel + +from artemis.config import ( + AgentNode, + AgentNodeWithFallback, + LLMUtilsNode, + LLMUtilsNodeWithFallback, + LLMWithFallback, + get_default_llm_config, + settings, +) +from artemis.context import ArtemisContext +from artemis.data_engine.engine import _CURRENT_DATA_ENGINE +from artemis.data_engine.trace import CURRENT_TRACE_ID, DataEngineCallbackHandler +from artemis.llm.router import ModelEndpoint, ModelFactory, ModelProvider +from artemis.utils.logger import get_logger + +# Logger for internal messages +llm_logger = logging.getLogger(__name__) +# Logger for user-facing messages +user_messages_logger = get_logger(__name__) + +T = TypeVar("T") + + +def _handle_llm_pause_and_resume(last_error: Exception) -> Path: + """Handles pausing task execution upon persistent failure until resumed.""" + err_msg = str(last_error) + if len(err_msg) > 1000: + err_msg = err_msg[:1000] + "... [Truncated by LLM Wrapper]" + + llm_logger.warning(f"LLM Error: {err_msg}. Pausing execution to wait for resume signal...") + pause_file = Path(settings.TRACES_PATH).parent / ".artemis_paused" + try: + pause_file.write_text(f"LLM Error: {err_msg}") + except Exception: + pass + + if _CURRENT_DATA_ENGINE: + _CURRENT_DATA_ENGINE._publish("task_paused", {"error": err_msg}) + + return pause_file + + +def robust_retry_async(func: Callable) -> Callable: + """Decorator to add robust retry and pause/resume logic to an async LLM call.""" + + @functools.wraps(func) + async def wrapper(*args, **kwargs): + while True: + max_retries = 3 + retry_delay = 0.0 + last_error = None + for attempt in range(max_retries): + try: + return await func(*args, **kwargs) + except Exception as e: + last_error = e + err_str = str(e).lower() + if attempt == max_retries - 1: + break + if attempt > 0: + retry_delay = 0.5 * (2 ** (attempt - 1)) + random.uniform(0.1, 0.3) + if any( + err in err_str + for err in [ + "503", + "throttled", + "overloaded", + "429", + "quota", + "resource_exhausted", + ] + ): + retry_delay = max(retry_delay, 10.0) + if retry_delay > 0: + llm_logger.warning(f"LLM call failed, retrying in {retry_delay:.2f}s...") + + if _CURRENT_DATA_ENGINE: + _CURRENT_DATA_ENGINE._publish( + "llm_retrying", + { + "attempt": attempt + 1, + "max_retries": max_retries, + "delay": retry_delay, + "error": str(e), + }, + ) + await asyncio.sleep(retry_delay) + + pause_file = _handle_llm_pause_and_resume(last_error) + while pause_file.exists(): + await asyncio.sleep(1) + + llm_logger.info("Resume signal received, retrying LLM call...") + + if _CURRENT_DATA_ENGINE: + _CURRENT_DATA_ENGINE._publish("task_resumed", {}) + + return wrapper + + +def robust_retry_astream(func: Callable) -> Callable: + """Decorator to add robust retry and pause/resume logic to an async stream LLM call.""" + + @functools.wraps(func) + async def wrapper(*args, **kwargs): + while True: + max_retries = 3 + retry_delay = 0.0 + last_error = None + for attempt in range(max_retries): + try: + gen = func(*args, **kwargs) + first = await gen.__anext__() + yield first + async for chunk in gen: + yield chunk + return + except Exception as e: + last_error = e + err_str = str(e).lower() + if attempt == max_retries - 1: + break + if attempt > 0: + retry_delay = 0.5 * (2 ** (attempt - 1)) + random.uniform(0.1, 0.3) + if any( + err in err_str + for err in [ + "503", + "throttled", + "overloaded", + "429", + "quota", + "resource_exhausted", + ] + ): + retry_delay = max(retry_delay, 10.0) + if retry_delay > 0: + llm_logger.warning( + f"LLM stream handshake throttled, retrying in {retry_delay:.2f}s..." + ) + + if _CURRENT_DATA_ENGINE: + _CURRENT_DATA_ENGINE._publish( + "llm_retrying", + { + "attempt": attempt + 1, + "max_retries": max_retries, + "delay": retry_delay, + "error": str(e), + }, + ) + await asyncio.sleep(retry_delay) + + pause_file = _handle_llm_pause_and_resume(last_error) + while pause_file.exists(): + await asyncio.sleep(1) + + llm_logger.info("Resume signal received, retrying LLM stream...") + + if _CURRENT_DATA_ENGINE: + _CURRENT_DATA_ENGINE._publish("task_resumed", {}) + + return wrapper + + +def _inject_parent_trace_id(trace_id, *args, **kwargs): + config = kwargs.get("config") or {} + if not kwargs.get("config") and len(args) > 1: + args_list = list(args) + if isinstance(args_list[1], dict): + config = args_list[1] + args_list[1] = config + args = tuple(args_list) + + if "metadata" not in config: + config["metadata"] = {} + config["metadata"]["parent_trace_id"] = str(trace_id) + kwargs["config"] = config + return args, kwargs + + +class RobustChatModelWrapper: + """Universal LangChain Runnable wrapper ensuring robust execution and telemetry.""" + + def __init__( + self, + base_model: BaseChatModel, + ctx: ArtemisContext = None, + endpoint: ModelEndpoint | None = None, + ): + self.base_model = base_model + self.ctx = ctx + self.endpoint = endpoint + + def __getattr__(self, item: str): + return getattr(self.base_model, item) + + def bind_tools(self, *args, **kwargs): + # Extract tools list from positional args or keyword args + tools_list = None + other_args = list(args) + if other_args: + tools_list = other_args.pop(0) + elif "tools" in kwargs: + tools_list = kwargs.pop("tools") + + processed_tools = list(tools_list) if tools_list is not None else [] + + # Determine if the underlying model provider is Google / Gemini + is_google_provider = False + if self.endpoint is not None: + is_google_provider = self.endpoint.provider in ( + ModelProvider.GOOGLE, + ModelProvider.GEMINI, + ) + else: + base_model_cls = getattr(self.base_model, "__class__", None) + base_model_name = base_model_cls.__name__ if base_model_cls else "" + is_google_provider = ( + "Google" in base_model_name + or "Gemini" in base_model_name + or "VertexAI" in base_model_name + ) + + has_explicit_google_search = any( + isinstance(t, dict) and "google_search" in t for t in processed_tools + ) + should_ground = ( + self.endpoint is not None and self.endpoint.enable_grounding + ) or has_explicit_google_search + + if is_google_provider: + if should_ground: + # Add native Google Search Grounding tool if not already present + if not has_explicit_google_search: + processed_tools.append({"google_search": {}}) + + # Ensure include_server_side_tool_invocations is True in tool_config + # to satisfy Gemini API requirements when mixing built-in tools and function calling + tool_config = kwargs.get("tool_config") + if tool_config is None: + kwargs["tool_config"] = {"include_server_side_tool_invocations": True} + elif isinstance(tool_config, dict): + tool_config = dict(tool_config) + tool_config.setdefault("include_server_side_tool_invocations", True) + kwargs["tool_config"] = tool_config + else: + # If not Gemini/Google (e.g. OpenAI, Anthropic, Ollama, OpenRouter, XAI): + # Ignore grounding and filter out any Gemini-specific built-in dict tools to avoid provider errors + processed_tools = [ + t for t in processed_tools if not (isinstance(t, dict) and "google_search" in t) + ] + + bound = self.base_model.bind_tools(processed_tools, *other_args, **kwargs) + return RobustChatModelWrapper(bound, self.ctx, endpoint=self.endpoint) + + def with_config(self, *args, **kwargs): + return RobustChatModelWrapper( + self.base_model.with_config(*args, **kwargs), + self.ctx, + endpoint=self.endpoint, + ) + + def with_structured_output(self, *args, **kwargs): + if hasattr(self.base_model, "with_structured_output"): + return RobustChatModelWrapper( + self.base_model.with_structured_output(*args, **kwargs), + self.ctx, + endpoint=self.endpoint, + ) + raise AttributeError(f"{type(self.base_model)} does not support with_structured_output") + + def __or__(self, other): + return RobustChatModelWrapper( + self.base_model.__or__(other), self.ctx, endpoint=self.endpoint + ) + + def __ror__(self, other): + return RobustChatModelWrapper( + self.base_model.__ror__(other), self.ctx, endpoint=self.endpoint + ) + + @robust_retry_async + async def ainvoke(self, *args, **kwargs): + trace_id = None + if self.ctx and self.ctx.data_engine: + trace_id = CURRENT_TRACE_ID.get() + if trace_id: + args, kwargs = _inject_parent_trace_id(trace_id, *args, **kwargs) + return await self.base_model.ainvoke(*args, **kwargs) + + @robust_retry_astream + async def astream(self, *args, **kwargs): + trace_id = None + if self.ctx and self.ctx.data_engine: + trace_id = CURRENT_TRACE_ID.get() + + if trace_id: + args, kwargs = _inject_parent_trace_id(trace_id, *args, **kwargs) + + stream_exec_id = uuid4() + + async for chunk in self.base_model.astream(*args, **kwargs): + if trace_id and self.ctx and self.ctx.data_engine and chunk.content: + text_to_stream = "" + thinking_to_stream = "" + if isinstance(chunk.content, str): + text_to_stream = chunk.content + elif isinstance(chunk.content, list): + for item in chunk.content: + if isinstance(item, dict): + if item.get("type") == "text": + text_to_stream += item.get("text", "") + elif item.get("type") == "thinking": + thinking_to_stream += item.get("thinking", "") + + if text_to_stream: + self.ctx.data_engine.stream_output( + stream_exec_id, text_to_stream, is_thinking=False + ) + if thinking_to_stream: + self.ctx.data_engine.stream_output( + stream_exec_id, thinking_to_stream, is_thinking=True + ) + yield chunk + + +async def invoke_llm_with_timeout_message[T]( + llm_call: Coroutine[Any, Any, T], + timeout_seconds: int = 10, + hard_timeout: int = 180, +) -> T: + """Send an LLM call and display a countdown / timeout message if delayed.""" + llm_task = asyncio.create_task(llm_call) + waiter_task = asyncio.create_task(asyncio.sleep(timeout_seconds)) + + done, _ = await asyncio.wait({llm_task, waiter_task}, return_when=asyncio.FIRST_COMPLETED) + + if llm_task in done: + waiter_task.cancel() + return llm_task.result() + else: + user_messages_logger.info("Waiting for LLM call response...") + start_time = asyncio.get_event_loop().time() + + while True: + try: + return await asyncio.wait_for(asyncio.shield(llm_task), timeout=1.0) + except TimeoutError: + if llm_task.done(): + return llm_task.result() + + pause_file = Path(settings.TRACES_PATH).parent / ".artemis_paused" + if pause_file.exists(): + start_time = asyncio.get_event_loop().time() + continue + + elapsed = asyncio.get_event_loop().time() - start_time + if elapsed > max(0, hard_timeout - timeout_seconds): + user_messages_logger.error(f"LLM call timed out after {hard_timeout} seconds.") + raise TimeoutError(f"LLM call timed out after {hard_timeout} seconds.") + + +# Backward compatible factory functions delegating to ModelFactory +def get_google_llm( + model_name: str = "gemini-3.6-flash", + temperature: float | None = None, + timeout: float | None = None, + thinking_budget: int | None = None, + thinking_level: str | None = "medium", + include_thoughts: bool | None = None, + enable_grounding: bool = False, +) -> BaseChatModel: + ep = ModelEndpoint( + provider=ModelProvider.GOOGLE, + model_name=model_name, + temperature=temperature or 0.0, + timeout_seconds=timeout or 60.0, + thinking_budget=thinking_budget, + thinking_level=thinking_level, + include_thoughts=include_thoughts, + enable_grounding=enable_grounding, + ) + return ModelFactory.create_model(ep) + + +def get_vertex_llm( + model_name: str = "gemini-3.6-flash", + temperature: float | None = None, + timeout: float | None = None, + thinking_budget: int | None = None, +) -> BaseChatModel: + ep = ModelEndpoint( + provider=ModelProvider.VERTEX_AI, + model_name=model_name, + temperature=temperature or 0.0, + timeout_seconds=timeout or 60.0, + thinking_budget=thinking_budget, + ) + return ModelFactory.create_model(ep) + + +def get_openai_llm( + model_name: str = "o3", + temperature: float | None = None, + timeout: float | None = None, +) -> BaseChatModel: + ep = ModelEndpoint( + provider=ModelProvider.OPENAI, + model_name=model_name, + temperature=temperature or 0.0, + timeout_seconds=timeout or 60.0, + ) + return ModelFactory.create_model(ep) + + +def get_openrouter_llm( + model_name: str, + temperature: float | None = None, + timeout: float | None = None, +) -> BaseChatModel: + ep = ModelEndpoint( + provider=ModelProvider.OPENROUTER, + model_name=model_name, + temperature=temperature or 0.0, + timeout_seconds=timeout or 60.0, + ) + return ModelFactory.create_model(ep) + + +def get_grok_llm( + model_name: str, + temperature: float | None = None, + timeout: float | None = None, +) -> BaseChatModel: + ep = ModelEndpoint( + provider=ModelProvider.XAI, + model_name=model_name, + temperature=temperature or 0.0, + timeout_seconds=timeout or 60.0, + ) + return ModelFactory.create_model(ep) + + +def get_anthropic_llm( + model_name: str, + temperature: float | None = None, + timeout: float | None = None, + thinking_budget: int | None = None, + reasoning_effort: str | None = None, +) -> BaseChatModel: + ep = ModelEndpoint( + provider=ModelProvider.ANTHROPIC, + model_name=model_name, + temperature=temperature or 0.0, + timeout_seconds=timeout or 60.0, + thinking_budget=thinking_budget, + reasoning_effort=reasoning_effort, + ) + return ModelFactory.create_model(ep) + + +def get_cached_raw_model( + provider: str, + model_name: str, + temperature: float | None = None, + timeout: float | None = None, + thinking_budget: int | None = None, + thinking_level: str | None = None, + include_thoughts: bool | None = None, + reasoning_effort: str | None = None, + enable_grounding: bool = False, +) -> BaseChatModel: + """Retrieves or instantiates a cached raw LangChain chat model.""" + ep = ModelEndpoint( + provider=ModelProvider.from_string(provider), + model_name=model_name, + temperature=temperature or 0.0, + timeout_seconds=timeout or 60.0, + thinking_budget=thinking_budget, + thinking_level=thinking_level, + include_thoughts=include_thoughts, + reasoning_effort=reasoning_effort, + enable_grounding=enable_grounding, + ) + return ModelFactory.get_model(ep) + + +def _resolve_endpoint( + ctx: ArtemisContext, + name: str, + is_utils: bool = False, + use_fallback: bool = False, +) -> ModelEndpoint: + """Cleanly resolves a ModelEndpoint from context router or llm_config.""" + # 1. Prefer dynamic ModelRouter if attached to context + if getattr(ctx, "model_router", None) is not None: + role_name = str(name).lower() + if use_fallback: + fallbacks = ctx.model_router.get_fallbacks(role_name) + if fallbacks: + return fallbacks[0] + return ctx.model_router.get_endpoint(role_name) + + # 2. Fall back to context llm_config + if getattr(ctx, "llm_config", None) is None: + try: + ctx.llm_config = get_default_llm_config() + except Exception: + pass + + cfg = ctx.llm_config.get_utils(name) if is_utils else ctx.llm_config.get_agent(name) + + if use_fallback: + if isinstance(cfg, LLMWithFallback) or (hasattr(cfg, "fallback") and cfg.fallback): + cfg = cfg.fallback + else: + raise ValueError(f"LLM configuration for '{name}' has no fallback!") + + def _get_val(obj, attr, expected_type): + val = getattr(obj, attr, None) + return val if isinstance(val, expected_type) else None + + provider_val = getattr(cfg, "provider", "google") + model_val = getattr(cfg, "model", "gemini-2.5-flash") + + return ModelEndpoint( + provider=ModelProvider.from_string(provider_val), + model_name=str(model_val), + temperature=_get_val(cfg, "temperature", (int, float)) or 0.0, + timeout_seconds=_get_val(cfg, "timeout", (int, float)) or 60.0, + thinking_budget=_get_val(cfg, "thinking_budget", int), + thinking_level=_get_val(cfg, "thinking_level", str), + reasoning_effort=_get_val(cfg, "reasoning_effort", str), + include_thoughts=_get_val(cfg, "include_thoughts", bool), + enable_grounding=_get_val(cfg, "enable_grounding", bool) or False, + ) + + +@overload +def get_llm( + ctx: ArtemisContext, + name: AgentNodeWithFallback, + *, + use_fallback: bool = False, + temperature: float | None = None, +) -> BaseChatModel: ... + + +@overload +def get_llm( + ctx: ArtemisContext, + name: LLMUtilsNode, + *, + is_utils: Literal[True], + temperature: float | None = None, +) -> BaseChatModel: ... + + +@overload +def get_llm( + ctx: ArtemisContext, + name: LLMUtilsNodeWithFallback, + *, + is_utils: Literal[True], + use_fallback: bool = False, + temperature: float | None = None, +) -> BaseChatModel: ... + + +def get_llm( + ctx: ArtemisContext, + name: AgentNode | LLMUtilsNode | AgentNodeWithFallback, + is_utils: bool = False, + use_fallback: bool = False, + temperature: float | None = None, +) -> BaseChatModel: + """Resolves and instantiates the appropriate LLM wrapper for the given agent role.""" + endpoint = _resolve_endpoint(ctx, str(name), is_utils=is_utils, use_fallback=use_fallback) + if temperature is not None: + endpoint = replace(endpoint, temperature=temperature) + raw_model = ModelFactory.get_model(endpoint) + + handler = DataEngineCallbackHandler(ctx) + bound_model = raw_model.with_config(callbacks=[handler]) + + return RobustChatModelWrapper(bound_model, ctx, endpoint=endpoint) # type: ignore + + +async def with_fallback[T]( + main_call: Callable[[], Awaitable[T]], + fallback_call: Callable[[], Awaitable[T]], + none_should_fallback: bool = True, +) -> T: + try: + result = await main_call() + if result is None and none_should_fallback: + llm_logger.warning("Main LLM inference returned None. Falling back...") + return await fallback_call() + return result + except Exception as e: + llm_logger.warning(f"โ— Main LLM inference failed: {e}. Falling back...") + return await fallback_call() diff --git a/artemis/telemetry/__init__.py b/artemis/telemetry/__init__.py new file mode 100644 index 0000000..6431741 --- /dev/null +++ b/artemis/telemetry/__init__.py @@ -0,0 +1,35 @@ +# 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. + +"""Telemetry and distributed tracing layer for ARTEMIS.""" + +from artemis.telemetry.models import ( + SpanType, + TelemetrySpan, + TokenUsage, +) +from artemis.telemetry.tracer import ( + CURRENT_PARENT_SPAN_ID, + CURRENT_TRACE_ID, + TelemetryTracer, +) + +__all__ = [ + "TelemetryTracer", + "TelemetrySpan", + "SpanType", + "TokenUsage", + "CURRENT_TRACE_ID", + "CURRENT_PARENT_SPAN_ID", +] diff --git a/artemis/telemetry/models.py b/artemis/telemetry/models.py new file mode 100644 index 0000000..919eab7 --- /dev/null +++ b/artemis/telemetry/models.py @@ -0,0 +1,64 @@ +# 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. + +"""Telemetry Data Models and Schema Definitions.""" + +from enum import Enum +import time +from typing import Any +from uuid import uuid4 +from pydantic import BaseModel, Field + + +class SpanType(str, Enum): + AGENT = "agent" + LLM_CALL = "llm_call" + TOOL_EXECUTION = "tool_execution" + DEVICE_ACTION = "device_action" + SCREEN_CAPTURE = "screen_capture" + CUSTOM = "custom" + + +class TokenUsage(BaseModel): + """LLM token consumption metrics.""" + + prompt_tokens: int = Field(default=0, description="Input prompt tokens") + completion_tokens: int = Field(default=0, description="Output generated tokens") + total_tokens: int = Field(default=0, description="Total tokens consumed") + + +class TelemetrySpan(BaseModel): + """Structured telemetry span representing a single unit of agentic execution.""" + + span_id: str = Field(default_factory=lambda: str(uuid4())) + trace_id: str = Field(..., description="Root trace/session identifier") + parent_id: str | None = Field(default=None, description="Parent span identifier") + name: str = Field(..., description="Human-readable span operation name") + type: SpanType = Field(default=SpanType.CUSTOM, description="Span categorization") + start_time: float = Field(default_factory=time.time, description="Start timestamp") + end_time: float | None = Field(default=None, description="Completion timestamp") + duration_ms: float | None = Field(default=None, description="Elapsed execution time in ms") + status: str = Field(default="running", description="Status ('running', 'success', 'failed')") + payload: dict[str, Any] = Field( + default_factory=dict, description="Custom parameters and payloads" + ) + error: str | None = Field(default=None, description="Exception message if failed") + token_usage: TokenUsage | None = Field(default=None, description="Token consumption metrics") + + def finish(self, status: str = "success", error: str | None = None) -> None: + """Closes the span and calculates duration.""" + self.end_time = time.time() + self.duration_ms = (self.end_time - self.start_time) * 1000.0 + self.status = status + self.error = error diff --git a/artemis/telemetry/replay.py b/artemis/telemetry/replay.py new file mode 100644 index 0000000..bc4581d --- /dev/null +++ b/artemis/telemetry/replay.py @@ -0,0 +1,34 @@ +# 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. + +"""Trace Replay and Inspection utilities.""" + +from pathlib import Path +from typing import Any +from artemis.telemetry.storage.json_storage import JsonStorage + + +class TraceReplayer: + """Loads and formats historic task traces for replay and visual inspection.""" + + def __init__(self, traces_dir: Path | str = "traces"): + self.storage = JsonStorage(storage_dir=traces_dir) + + def get_trace_summary(self, trace_id: str) -> dict[str, Any]: + spans = self.storage.read_spans(trace_id) + return { + "trace_id": trace_id, + "total_spans": len(spans), + "spans": spans, + } diff --git a/artemis/telemetry/storage/json_storage.py b/artemis/telemetry/storage/json_storage.py new file mode 100644 index 0000000..a928e54 --- /dev/null +++ b/artemis/telemetry/storage/json_storage.py @@ -0,0 +1,47 @@ +# 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. + +"""JSON Lines Telemetry Storage Backend.""" + +import json +from pathlib import Path +from typing import Any +from artemis.telemetry.models import TelemetrySpan + + +class JsonStorage: + """Persists spans into newline-delimited JSON files.""" + + def __init__(self, storage_dir: Path | str = "traces"): + self.storage_dir = Path(storage_dir) + + def write_span(self, trace_id: str, span: TelemetrySpan) -> None: + target_dir = self.storage_dir / trace_id + target_dir.mkdir(parents=True, exist_ok=True) + log_file = target_dir / "spans.jsonl" + + with open(log_file, "a", encoding="utf-8") as f: + f.write(json.dumps(span.model_dump(), ensure_ascii=False) + "\n") + + def read_spans(self, trace_id: str) -> list[dict[str, Any]]: + log_file = self.storage_dir / trace_id / "spans.jsonl" + if not log_file.exists(): + return [] + + spans = [] + with open(log_file, encoding="utf-8") as f: + for line in f: + if line.strip(): + spans.append(json.loads(line)) + return spans diff --git a/artemis/telemetry/tracer.py b/artemis/telemetry/tracer.py new file mode 100644 index 0000000..241dec7 --- /dev/null +++ b/artemis/telemetry/tracer.py @@ -0,0 +1,85 @@ +# 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. + +"""Telemetry Tracer and Span Life Cycle Manager.""" + +from contextlib import asynccontextmanager +from contextvars import ContextVar +import json +from pathlib import Path +from typing import Any +from collections.abc import AsyncIterator +from uuid import uuid4 + +from artemis.telemetry.models import SpanType, TelemetrySpan +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + +CURRENT_TRACE_ID: ContextVar[str] = ContextVar("CURRENT_TRACE_ID", default="global-session") +CURRENT_PARENT_SPAN_ID: ContextVar[str | None] = ContextVar("CURRENT_PARENT_SPAN_ID", default=None) + + +class TelemetryTracer: + """Central tracer for recording and persisting telemetry spans.""" + + def __init__(self, trace_id: str | None = None, traces_dir: Path | None = None): + self.trace_id = trace_id or str(uuid4()) + self.traces_dir = traces_dir or Path("traces") / self.trace_id + self.spans: list[TelemetrySpan] = [] + self._log_file: Path | None = None + + def _ensure_log_file(self) -> Path: + if self._log_file is None: + self.traces_dir.mkdir(parents=True, exist_ok=True) + self._log_file = self.traces_dir / "telemetry_spans.jsonl" + return self._log_file + + def record_span(self, span: TelemetrySpan) -> None: + """Appends and persists completed span.""" + self.spans.append(span) + try: + log_file = self._ensure_log_file() + with open(log_file, "a", encoding="utf-8") as f: + f.write(json.dumps(span.model_dump(), ensure_ascii=False) + "\n") + except Exception as e: + logger.warning(f"Failed to persist telemetry span {span.span_id}: {e}") + + @asynccontextmanager + async def start_span( + self, + name: str, + span_type: SpanType = SpanType.CUSTOM, + payload: dict[str, Any] | None = None, + ) -> AsyncIterator[TelemetrySpan]: + """Asynchronous context manager creating and recording a span.""" + parent_id = CURRENT_PARENT_SPAN_ID.get() + span = TelemetrySpan( + trace_id=self.trace_id, + parent_id=parent_id, + name=name, + type=span_type, + payload=payload or {}, + ) + + token = CURRENT_PARENT_SPAN_ID.set(span.span_id) + try: + yield span + span.finish(status="success") + except Exception as e: + span.finish(status="failed", error=str(e)) + raise + finally: + CURRENT_PARENT_SPAN_ID.reset(token) + self.record_span(span) diff --git a/artemis/tools/actions/__init__.py b/artemis/tools/actions/__init__.py new file mode 100644 index 0000000..f6db1a6 --- /dev/null +++ b/artemis/tools/actions/__init__.py @@ -0,0 +1,45 @@ +# 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. + +"""Universal device action tools.""" + +from artemis.tools.actions.device_actions import ( + ClickArgs, + InputTextArgs, + LongPressArgs, + PressKeyArgs, + SwipeArgs, + WaitForDelayArgs, + click, + input_text, + long_press, + press_key, + swipe, + wait_for_delay, +) + +__all__ = [ + "click", + "long_press", + "input_text", + "swipe", + "press_key", + "wait_for_delay", + "ClickArgs", + "LongPressArgs", + "InputTextArgs", + "SwipeArgs", + "PressKeyArgs", + "WaitForDelayArgs", +] diff --git a/artemis/tools/actions/device_actions.py b/artemis/tools/actions/device_actions.py new file mode 100644 index 0000000..01540d6 --- /dev/null +++ b/artemis/tools/actions/device_actions.py @@ -0,0 +1,237 @@ +# 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. + +"""Universal device action tools implemented on top of BaseDeviceDriver.""" + +import asyncio +from typing import Any, Literal +from artemis.drivers.base import BaseDeviceDriver +from artemis.tools.base import artemis_tool +from artemis.utils.coordinates import parse_swipe_parameters +from artemis.utils.logger import get_logger +from pydantic import BaseModel, Field + +logger = get_logger(__name__) + + +# Schemas +class ClickArgs(BaseModel): + target: list[int] = Field( + ..., description="Click target normalized coordinates [x, y] in 0-1000 scale." + ) + + +class LongPressArgs(BaseModel): + target: list[int] = Field( + ..., description="Long press target normalized coordinates [x, y] in 0-1000 scale." + ) + duration: int = Field(1000, description="Long press duration in milliseconds (default 1000).") + + +class InputTextArgs(BaseModel): + text: str = Field(..., description="The text content to input.") + target: list[int] = Field( + ..., description="Input target field normalized coordinates [x, y] in 0-1000 scale." + ) + clear_exist: bool = Field( + True, description="Whether to clear existing text in the input box before typing." + ) + + +class SwipeArgs(BaseModel): + direction: Literal["up", "down", "left", "right"] | None = Field( + None, + description=( + "Smart directional swipe ('up', 'down', 'left', 'right') for general page" + " scrolling/browsing. Automatically applies safe vectors and a ~40% visual" + " overlap anchor: 'up' (drags bottom-to-top, scrolling down to reveal content below)," + " 'down' (drags top-to-bottom, scrolling up to reveal content above)," + " 'left' (drags right-to-left, scrolling right)," + " 'right' (drags left-to-right, scrolling left)." + ), + ) + start: list[int] | None = Field( + None, + description=( + "Start normalized coordinates [start_x, start_y] in 0-1000 scale for precise," + " local interactions (e.g. adjusting sliders, SeekBars, fine range selection, or drag-and-drop)." + ), + ) + end: list[int] | None = Field( + None, + description=( + "End normalized coordinates [end_x, end_y] in 0-1000 scale for precise," + " local interactions (e.g. adjusting sliders, SeekBars, fine range selection, or drag-and-drop)." + ), + ) + action: Literal["up", "down", "left", "right"] | list[int] | None = Field( + None, + description=( + "Backward-compatible swipe gesture: smart direction string ('up', 'down', 'left', 'right')" + " OR precise custom coordinates [start_x, start_y, end_x, end_y] in 0-1000 scale." + ), + ) + duration: int = Field( + 400, + description=( + "Swipe/drag duration in milliseconds (default 400). For drag-and-drop, " + "list reordering, or sliding/adjusting sliders (e.g. volume, brightness, SeekBars), " + "set duration >= 1000 (e.g. 1500)." + ), + ) + + +class PressKeyArgs(BaseModel): + key: str = Field( + ..., + description="Key name to press (e.g. 'home', 'back', 'enter', 'delete', 'power').", + ) + + +class WaitForDelayArgs(BaseModel): + seconds: float = Field( + ..., description="Duration in seconds to wait for animations, pages, or state updates." + ) + + +# Universal Tool Definitions +@artemis_tool( + name="click", + description="[ACTION] Clicks on the specified screen coordinate [x, y] in normalized 0-1000 scale.", + args_schema=ClickArgs, + category="action", +) +async def click(target: list[int], driver: BaseDeviceDriver) -> str: + width, height = driver.screen_size + abs_x = int(target[0] * width / 1000.0) + abs_y = int(target[1] * height / 1000.0) + success = await driver.tap(abs_x, abs_y) + return ( + f"Click executed at normalized {target} (pixels: {abs_x}, {abs_y})" + if success + else f"Click failed at {target}" + ) + + +@artemis_tool( + name="long_press", + description="[ACTION] Long presses on the specified screen coordinate [x, y] in normalized 0-1000 scale.", + args_schema=LongPressArgs, + category="action", +) +async def long_press( + target: list[int], duration: int = 1000, driver: BaseDeviceDriver = None +) -> str: + width, height = driver.screen_size + abs_x = int(target[0] * width / 1000.0) + abs_y = int(target[1] * height / 1000.0) + success = await driver.long_press(abs_x, abs_y, duration_ms=duration) + return ( + f"Long press executed at normalized {target} for {duration}ms" + if success + else f"Long press failed at {target}" + ) + + +@artemis_tool( + name="input_text", + description="[ACTION] Clicks an input field target [x, y] and types text into it.", + args_schema=InputTextArgs, + category="action", +) +async def input_text( + text: str, target: list[int], clear_exist: bool = True, driver: BaseDeviceDriver = None +) -> str: + # First tap to focus + width, height = driver.screen_size + abs_x = int(target[0] * width / 1000.0) + abs_y = int(target[1] * height / 1000.0) + await driver.tap(abs_x, abs_y) + await asyncio.sleep(0.2) + + # Then input text + success = await driver.input_text(text, clear_existing=clear_exist) + return ( + f"Input text '{text}' into target {target}" + if success + else f"Failed to input text into {target}" + ) + + +@artemis_tool( + name="swipe", + description=( + "[ACTION] Performs smart directional swipe for browsing, or precise coordinate drag/slider gestures." + ), + args_schema=SwipeArgs, + category="action", +) +async def swipe( + action: Any = None, + direction: str | None = None, + start: list[int] | None = None, + end: list[int] | None = None, + duration: int = 400, + driver: BaseDeviceDriver = None, + **kwargs: Any, +) -> str: + params = dict(kwargs) + if action is not None: + params["action"] = action + if direction is not None: + params["direction"] = direction + if start is not None: + params["start"] = start + if end is not None: + params["end"] = end + params["duration"] = duration + + kind, target, dur = parse_swipe_parameters(params, default_duration=duration) + + if kind == "direction" and isinstance(target, str): + success = await driver.swipe_direction(target, duration_ms=dur) + return f"Swiped {target}" if success else f"Swipe {target} failed" + elif kind == "coords" and isinstance(target, list) and len(target) == 4: + width, height = driver.screen_size + start_x = int(target[0] * width / 1000.0) + start_y = int(target[1] * height / 1000.0) + end_x = int(target[2] * width / 1000.0) + end_y = int(target[3] * height / 1000.0) + success = await driver.swipe(start_x, start_y, end_x, end_y, duration_ms=dur) + return ( + f"Swiped from {target[:2]} to {target[2:]}" if success else "Swipe coordinates failed" + ) + return f"Invalid swipe action parameter: {action or direction or (start, end)}" + + +@artemis_tool( + name="press_key", + description="[ACTION] Presses a hardware or system key on the device ('home', 'back', 'enter', 'delete', 'power').", + args_schema=PressKeyArgs, + category="action", +) +async def press_key(key: str, driver: BaseDeviceDriver) -> str: + success = await driver.press_key(key) + return f"Key '{key}' pressed" if success else f"Failed to press key '{key}'" + + +@artemis_tool( + name="wait_for_delay", + description="[ACTION] Pauses execution for a specified number of seconds to allow UI loading or transition.", + args_schema=WaitForDelayArgs, + category="action", +) +async def wait_for_delay(seconds: float, driver: BaseDeviceDriver = None) -> str: + await asyncio.sleep(seconds) + return f"Waited for {seconds} seconds." diff --git a/artemis/tools/base.py b/artemis/tools/base.py new file mode 100644 index 0000000..9540dac --- /dev/null +++ b/artemis/tools/base.py @@ -0,0 +1,274 @@ +# 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. + +"""Universal Tool Protocol and Registry for ARTEMIS. + +Enables tools to be authored once using standard Pydantic schemas and async handlers, +then automatically exported to LangChain Tools (Pro Graph), Google GenAI FunctionDeclarations +(FlashRunner), and MCP Tools (FastMCP) with zero code duplication. +""" + +import asyncio +from collections.abc import Callable +import concurrent.futures +import inspect +from typing import Annotated, Any, Literal + +from google.genai import types as genai_types +from langchain_core.tools import BaseTool, StructuredTool +from langgraph.prebuilt import InjectedState +from pydantic import BaseModel + +from artemis.drivers.base import BaseDeviceDriver +from artemis.drivers.factory import get_driver +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + +ToolCategory = Literal["action", "perception", "system", "memory", "custom"] + + +class ArtemisTool: + """Unified tool wrapper encapsulating schema, execution logic, and multi-protocol export.""" + + # pylint: disable=too-many-arguments,too-many-positional-arguments + def __init__( + self, + name: str, + description: str, + args_schema: type[BaseModel], + handler: Callable[..., Any] | None = None, + category: ToolCategory = "action", + ): + self.name = name + self.description = description + self.args_schema = args_schema + self.handler = handler + self.category = category + + async def execute(self, driver: BaseDeviceDriver, ctx: Any, **kwargs: Any) -> Any: + """Executes the tool handler with injected driver and context.""" + # Inspect handler signature to only pass supported arguments + if self.handler is not None: + sig = inspect.signature(self.handler) + call_kwargs = {} + for param_name in sig.parameters: + if param_name == "driver": + call_kwargs["driver"] = driver + elif param_name in ("ctx", "context"): + call_kwargs[param_name] = ctx + elif param_name in kwargs: + call_kwargs[param_name] = kwargs[param_name] + + if inspect.iscoroutinefunction(self.handler): + return await self.handler(**call_kwargs) + return self.handler(**call_kwargs) + raise NotImplementedError("Subclasses must implement execute() or provide a handler.") + + async def __call__(self, *args: Any, **kwargs: Any) -> Any: + """Allows ArtemisTool instances to be invoked directly as callables.""" + driver = kwargs.pop("driver", None) + ctx = kwargs.pop("ctx", None) or kwargs.pop("context", None) + return await self.execute(driver=driver, ctx=ctx, *args, **kwargs) + + def to_langchain_tool(self, ctx: Any, name: str | None = None) -> BaseTool: + """Exports this tool to a LangChain StructuredTool.""" + + def _get_context_driver() -> BaseDeviceDriver | None: + if ( + ctx is not None + and hasattr(ctx, "device") + and ctx.device + and getattr(ctx.device, "device_id", None) + ): + try: + return get_driver(ctx) + except Exception as e: # pylint: disable=broad-exception-caught + logger.debug(f"Failed to get driver: {e}") + return None + return getattr(ctx, "_active_driver", None) if ctx is not None else None + + def _func(state: Annotated[Any, InjectedState] = None, **kwargs): + driver = _get_context_driver() + if state is not None: + kwargs["state"] = state + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = None + + if loop and loop.is_running(): + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: + return pool.submit( + asyncio.run, self.execute(driver=driver, ctx=ctx, **kwargs) + ).result() + return asyncio.run(self.execute(driver=driver, ctx=ctx, **kwargs)) + + async def _coroutine(state: Annotated[Any, InjectedState] = None, **kwargs): + driver = _get_context_driver() + if state is not None: + kwargs["state"] = state + return await self.execute(driver=driver, ctx=ctx, **kwargs) + + return StructuredTool.from_function( + func=_func, + coroutine=_coroutine, + name=name or self.name, + description=self.description, + args_schema=self.args_schema, + ) + + def to_genai_declaration(self) -> genai_types.FunctionDeclaration: + """Exports this tool to a Google GenAI FunctionDeclaration schema.""" + properties = {} + required = [] + + if self.args_schema: + schema = self.args_schema.model_json_schema() + props = schema.get("properties", {}) + required = schema.get("required", []) + + for prop_name, prop_data in props.items(): + p_type_str = prop_data.get("type") + items_data = prop_data.get("items") + if not p_type_str and "anyOf" in prop_data: + for sub_schema in prop_data["anyOf"]: + sub_type = sub_schema.get("type") + if sub_type and sub_type != "null": + p_type_str = sub_type + if "items" in sub_schema and not items_data: + items_data = sub_schema["items"] + break + if not p_type_str: + p_type_str = "string" + + p_desc = prop_data.get("description", "") + + genai_type = genai_types.Type.STRING + if p_type_str == "integer": + genai_type = genai_types.Type.INTEGER + elif p_type_str == "number": + genai_type = genai_types.Type.NUMBER + elif p_type_str == "boolean": + genai_type = genai_types.Type.BOOLEAN + elif p_type_str == "array": + genai_type = genai_types.Type.ARRAY + elif p_type_str == "object": + genai_type = genai_types.Type.OBJECT + + # Handle array items if present + items_schema = None + if genai_type == genai_types.Type.ARRAY: + item_type_str = (items_data or {}).get("type", "string") + item_genai_type = ( + genai_types.Type.INTEGER + if item_type_str == "integer" + else genai_types.Type.STRING + ) + items_schema = genai_types.Schema(type=item_genai_type) + + properties[prop_name] = genai_types.Schema( + type=genai_type, + description=p_desc, + items=items_schema, + ) + + return genai_types.FunctionDeclaration( + name=self.name, + description=self.description, + parameters=genai_types.Schema( + type=genai_types.Type.OBJECT, + properties=properties, + required=required, + ), + ) + + +class ToolRegistry: + """Global registry managing universal tools.""" + + _tools: dict[str, ArtemisTool] = {} + + @classmethod + def register(cls, tool: ArtemisTool) -> ArtemisTool: + """Registers a tool in the global registry.""" + cls._tools[tool.name] = tool + return tool + + @classmethod + def get(cls, name: str) -> ArtemisTool | None: + """Retrieves a tool by name from the registry.""" + normalized = name.split(":")[-1] if ":" in name else name + return cls._tools.get(normalized) + + @classmethod + def list_tools(cls, category: ToolCategory | None = None) -> list[ArtemisTool]: + """Lists all registered tools, optionally filtered by category.""" + if category: + return [t for t in cls._tools.values() if t.category == category] + return list(cls._tools.values()) + + @classmethod + def get_langchain_tools(cls, ctx: Any, names: list[str] | None = None) -> list[BaseTool]: + """Returns LangChain tool wrappers for the requested or all registered tools.""" + tools = ( + [cls._tools[n] for n in names if n in cls._tools] + if names + else list(cls._tools.values()) + ) + return [t.to_langchain_tool(ctx) for t in tools] + + @classmethod + def get_genai_declarations( + cls, names: list[str] | None = None + ) -> list[genai_types.FunctionDeclaration]: + """Returns Google GenAI function declarations for registered tools.""" + tools = ( + [cls._tools[n] for n in names if n in cls._tools] + if names + else list(cls._tools.values()) + ) + return [t.to_genai_declaration() for t in tools] + + @classmethod + async def execute( + cls, name: str, args: dict[str, Any], driver: BaseDeviceDriver, ctx: Any + ) -> Any: + """Executes a registered tool by name with provided arguments, driver, and context.""" + tool = cls.get(name) + if not tool: + raise ValueError(f"Tool '{name}' not found in ToolRegistry.") + return await tool.execute(driver=driver, ctx=ctx, **args) + + +def artemis_tool( + name: str, + description: str, + args_schema: type[BaseModel], + category: ToolCategory = "action", +): + """Decorator to define and register a universal ARTEMIS tool.""" + + def decorator(fn: Callable[..., Any]) -> ArtemisTool: + tool_obj = ArtemisTool( + name=name, + description=description, + args_schema=args_schema, + handler=fn, + category=category, + ) + ToolRegistry.register(tool_obj) + return tool_obj + + return decorator diff --git a/artemis/tools/command_tool.py b/artemis/tools/command_tool.py new file mode 100644 index 0000000..c8ce3e9 --- /dev/null +++ b/artemis/tools/command_tool.py @@ -0,0 +1,836 @@ +# 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. + +"""Universal ADB command and background task management tools for ARTEMIS.""" + +import asyncio +from asyncio.subprocess import Process as AsyncProcess +import os +from typing import Any, Literal +import uuid + +from langchain_core.tools import BaseTool +from pydantic import BaseModel, Field + +from artemis.agents.log_analyzer.output_analyzer import TaskOutputAnalyzerNode +from artemis.context import ArtemisContext +from artemis.controllers.platform_specific_commands_controller import ( + get_adb_device, +) +from artemis.data_engine.trace import trace_langchain_tool +from artemis.drivers.base import BaseDeviceDriver +from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.tool_wrapper import ToolWrapper +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + +# Global persistent Android shell environments +_PERSISTENT_ENVIRONMENTS: dict[str, dict[str, str]] = {} +# Global background tasks +_BACKGROUND_TASKS: dict[str, "BackgroundTask"] = {} +# Finished tasks history logs cache (to support Action='analyze' even after tasks are unregistered) +_FINISHED_TASKS_LOGS: dict[str, dict[str, Any]] = {} + + +# pylint: disable=too-many-arguments,too-many-positional-arguments +def _register_finished_task( + task_id: str, + command: str, + cwd: str | None, + terminal_id: str | None, + status: str, + exit_code: int | None, + output: str, +): + """Caches execution details for finished tasks.""" + _FINISHED_TASKS_LOGS[task_id] = { + "task_id": task_id, + "command": command, + "cwd": cwd, + "terminal_id": terminal_id, + "status": status, + "exit_code": exit_code, + "output": output, + "notified": False, + } + # Keep memory bounded + if len(_FINISHED_TASKS_LOGS) > 10: + oldest_key = next(iter(_FINISHED_TASKS_LOGS)) + del _FINISHED_TASKS_LOGS[oldest_key] + + +def _is_output_long(output: str) -> bool: + """Checks if command output exceeds the max threshold for inline display.""" + return len(output.splitlines()) > 500 or len(output) > 50000 + + +def _format_long_output_response(_task_id: str, output: str, intro: str) -> str: + """Formats truncated output with guidance to use the analyzer tool.""" + lines = output.splitlines() + last_200 = "\n".join(lines[-200:]) + return ( + f"{intro}\nWarning: Output is too long ({len(lines)} lines," + f" {len(output)} chars) and has been truncated.\n--- Last 200 lines of" + f" output ---\n{last_200}\n------------------------\nYou can use the" + " 'analyze_task_output' tool to get more information." + ) + + +def _get_task_info(task_id: str, ctx: ArtemisContext) -> dict[str, Any] | None: + """Retrieves metadata and logs for a running or finished task.""" + if task_id in _BACKGROUND_TASKS: + t = _BACKGROUND_TASKS[task_id] + return { + "command": t.command, + "cwd": t.cwd, + "terminal_id": t.terminal_id, + "status": t.status, + "output": "".join(t.stdout_log), + } + if task_id in _FINISHED_TASKS_LOGS: + return _FINISHED_TASKS_LOGS[task_id] + if ctx and ctx.data_engine: + db_tasks = ctx.data_engine.get_all_background_tasks() + for db_t in db_tasks: + if db_t.get("task_id") == task_id: + return { + "command": db_t.get("summary") or "", + "cwd": None, + "terminal_id": None, + "status": db_t.get("status") or "completed", + "output": db_t.get("logs") or "", + } + return None + + +# pylint: disable=too-many-instance-attributes +class BackgroundTask: + """Represents an active or completing background ADB shell task.""" + + # pylint: disable=too-many-arguments,too-many-positional-arguments + def __init__( + self, + task_id: str, + command: str, + process: AsyncProcess, + terminal_id: str | None = None, + cwd: str | None = None, + ): + self.task_id = task_id + self.command = command + self.process = process + self.terminal_id = terminal_id + self.cwd = cwd + self.stdout_log: list[str] = [] + self.status = "running" + self.exit_code: int | None = None + self.trace_id: uuid.UUID | None = None + self.listener_task: asyncio.Task | None = None + + async def start(self, ctx: ArtemisContext): + """Starts the background output listener task.""" + self.listener_task = asyncio.create_task(self._listen(ctx)) + + async def _listen(self, ctx: ArtemisContext): + """Reads process output streams and updates lifecycle status.""" + try: + # Read stdout/stderr merged + while True: + line = await self.process.stdout.readline() + if not line: + break + decoded = line.decode() + self.stdout_log.append(decoded) + if ctx and ctx.data_engine and getattr(self, "trace_id", None): + ctx.data_engine.stream_output(self.trace_id, decoded) + + self.exit_code = await self.process.wait() + self.status = "completed" if self.exit_code == 0 else "failed" + logger.info( + f"Background ADB task {self.task_id} completed with exit code {self.exit_code}" + ) + + # Post-process output to extract env if persistent + output_text = "".join(self.stdout_log) + _, parsed_env = self.parse_output(output_text) + + if self.terminal_id and parsed_env: + _PERSISTENT_ENVIRONMENTS[self.terminal_id] = parsed_env + logger.info( + "Updated persistent env for Android terminal" + f" '{self.terminal_id}' from background task." + ) + + # Update status in DataEngine + if ctx and ctx.data_engine: + ctx.data_engine.unregister_background_task( + task_id=self.task_id, + status=self.status, + ) + except Exception as e: # pylint: disable=broad-exception-caught + logger.error(f"Error in background task {self.task_id} listener: {e}") + self.status = "error" + if ctx and ctx.data_engine: + ctx.data_engine.unregister_background_task(self.task_id, "failed") + finally: + # Cache logs in memory before removing task from active tasks + full_output = "".join(self.stdout_log) + _register_finished_task( + task_id=self.task_id, + command=self.command, + cwd=self.cwd, + terminal_id=self.terminal_id, + status=self.status, + exit_code=self.exit_code, + output=full_output, + ) + if self.task_id in _BACKGROUND_TASKS: + del _BACKGROUND_TASKS[self.task_id] + + @staticmethod + def parse_output(output: str) -> tuple[str, dict[str, str] | None]: + """Separates stdout lines from exported environment variables.""" + clean_lines = [] + env_lines = [] + in_env = False + + for line in output.splitlines(): + if "===EXIT_CODE===" in line: + continue + if "===ENV_START===" in line: + in_env = True + continue + if in_env: + env_lines.append(line) + else: + clean_lines.append(line) + + clean_output = "\n".join(clean_lines) + + parsed_env = None + if env_lines: + parsed_env = {} + for el in env_lines: + if "=" in el: + k, v = el.split("=", 1) + parsed_env[k] = v + + return clean_output, parsed_env + + # Backwards compatibility alias + _parse_output = parse_output + + +class RunAdbCommandArgs(BaseModel): + """Arguments schema for executing ADB shell commands.""" + + CommandLine: str = Field( + ..., + description=( + "The exact shell command line string to execute inside the Android device shell." + ), + ) + Cwd: str = Field( + default="/data/local/tmp", + description="The directory on the Android device to run the command in.", + ) + RunPersistent: bool = Field( + default=False, + description=( + "Set to true to run this command in a persistent terminal session" + " that preserves environment variables between invocations on the" + " phone." + ), + ) + RequestedTerminalID: str | None = Field( + default=None, + description=( + "Optional ID of a persistent terminal to reuse. Reuses environment" + " variables if provided." + ), + ) + WaitMsBeforeAsync: int = Field( + default=5000, + description=( + "Number of milliseconds to wait for the command to finish" + " synchronously before sending it to the background." + ), + ) + + +class RunAdbCommandTool(ArtemisTool): + """Universal tool for executing ADB shell commands on an Android device.""" + + def __init__(self): + super().__init__( + name="run_adb_command", + description=( + "[SHELL] Executes a shell command directly on the Android mobile " + "device via ADB shell. Can run synchronously or transition to a " + "background task if execution takes longer than WaitMsBeforeAsync. " + "Supports persistent environments (environment variables) on the " + "phone across invocations." + ), + args_schema=RunAdbCommandArgs, + category="system", + ) + + # pylint: disable=too-many-arguments,too-many-locals,too-many-branches,too-many-statements,too-many-return-statements,too-many-positional-arguments,too-many-boolean-expressions + async def execute( + self, + driver: BaseDeviceDriver | None = None, + ctx: ArtemisContext | None = None, + CommandLine: str | None = None, + Cwd: str = "/data/local/tmp", + RunPersistent: bool = False, + RequestedTerminalID: str | None = None, + WaitMsBeforeAsync: int = 5000, + **kwargs: Any, + ) -> str: + # Extract command line supporting alternative key formats + cmd_line = ( + CommandLine + if CommandLine is not None + else (kwargs.get("command_line") or kwargs.get("command") or kwargs.get("cmd") or "") + ) + cwd = Cwd or kwargs.get("cwd") or "/data/local/tmp" + run_persistent = ( + RunPersistent if RunPersistent is not False else kwargs.get("run_persistent", False) + ) + requested_terminal_id = ( + RequestedTerminalID + if RequestedTerminalID is not None + else kwargs.get("requested_terminal_id") + ) + wait_ms = ( + WaitMsBeforeAsync + if WaitMsBeforeAsync != 5000 + else kwargs.get("wait_ms_before_async", WaitMsBeforeAsync) + ) + + # Fallback to driver's execute_shell if no full ArtemisContext is available + if ( + (ctx is None or not hasattr(ctx, "device") or not ctx.device) + and driver is not None + and hasattr(driver, "execute_shell") + and not run_persistent + ): + return await driver.execute_shell(cmd_line) + + device_id = None + if ctx and hasattr(ctx, "device") and ctx.device: + device_id = getattr(ctx.device, "device_id", None) + elif driver and hasattr(driver, "device_id"): + device_id = getattr(driver, "device_id", None) + if not device_id: + device_id = "default_device" + + # Prepare environment + terminal_id = None + run_env = os.environ.copy() # Local host env to run adb command + android_env_vars = {} + + if run_persistent: + terminal_id = requested_terminal_id or f"term_{uuid.uuid4().hex[:8]}" + if terminal_id in _PERSISTENT_ENVIRONMENTS: + android_env_vars = _PERSISTENT_ENVIRONMENTS[terminal_id] + else: + _PERSISTENT_ENVIRONMENTS[terminal_id] = android_env_vars + + # Construct command to run inside adb shell + # We inject the cached environment variables by prepending 'export K=V' declarations + env_prepends = "" + for k, v in android_env_vars.items(): + env_prepends += f"export {k}='{v}'\n" + + # Construct final command script + phone_script = f"cd '{cwd}'\n" + env_prepends + f"{cmd_line}\n" + if run_persistent: + phone_script += ( + '_exit_code=$?\necho "===EXIT_CODE===$_exit_code"\necho "===ENV_START==="\nenv' + ) + + # In Cloud Mode or when adb client is virtualized + adb_client = getattr(ctx, "adb_client", None) if ctx else None + if os.environ.get("ARTEMIS_CLOUD_MODE") == "1" or ( + adb_client is not None and hasattr(adb_client, "_bridge") + ): + try: + device = get_adb_device(ctx) + output = device.shell(phone_script) + clean_output, new_env = BackgroundTask.parse_output(output) + if run_persistent and new_env: + android_env_vars.update(new_env) + return ( + clean_output + if clean_output.strip() + else "Command executed successfully with empty output." + ) + except Exception as e: # pylint: disable=broad-exception-caught + logger.error(f"Failed to execute virtual adb shell for '{cmd_line}': {e}") + return f"Failed to execute adb command: {e}" + + # We spawn the host subprocess: adb -s shell + # Pass phone_script as an argument so it doesn't hold the shell open if we keep stdin open. + try: + process = await asyncio.create_subprocess_exec( + "adb", + "-s", + device_id, + "shell", + phone_script, + cwd=os.getcwd(), + env=run_env, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, + stdin=asyncio.subprocess.PIPE, + ) + + except Exception as e: # pylint: disable=broad-exception-caught + logger.error(f"Failed to spawn adb shell process for command '{cmd_line}': {e}") + return f"Failed to spawn adb process: {e}" + + # Wait for wait_ms milliseconds + wait_seconds = wait_ms / 1000.0 + try: + stdout_data = [] + + async def read_limit(): + while True: + line = await process.stdout.readline() + if not line: + break + stdout_data.append(line.decode()) + await process.wait() + + await asyncio.wait_for(read_limit(), timeout=wait_seconds) + + # If we reached here, the process completed within the wait window + exit_code = process.returncode + output_text = "".join(stdout_data) + + # Clean output and extract environment + clean_output = output_text + parsed_env = None + if run_persistent: + clean_output, parsed_env = BackgroundTask.parse_output(output_text) + if parsed_env: + _PERSISTENT_ENVIRONMENTS[terminal_id] = parsed_env + + intro = f"ADB command completed with exit code {exit_code}." + if run_persistent: + intro += f" TerminalID: {terminal_id}." + + if _is_output_long(clean_output): + task_id = f"task_sync_{uuid.uuid4().hex[:8]}" + intro += f" TaskId: {task_id}." + _register_finished_task( + task_id=task_id, + command=cmd_line, + cwd=cwd, + terminal_id=terminal_id, + status="completed" if exit_code == 0 else "failed", + exit_code=exit_code, + output=clean_output, + ) + return _format_long_output_response(task_id, clean_output, intro) + + return f"{intro}\nOutput:\n{clean_output}" + + except TimeoutError: + # Process did not finish in time, move to background + task_id = f"task_{process.pid}" + trace_id = uuid.uuid4() + + bg_task = BackgroundTask( + task_id=task_id, + command=cmd_line, + process=process, + terminal_id=terminal_id, + cwd=cwd, + ) + bg_task.stdout_log = list(stdout_data) + bg_task.trace_id = trace_id + _BACKGROUND_TASKS[task_id] = bg_task + + # Register in DataEngine + if ctx and ctx.data_engine: + ctx.data_engine.register_background_task( + task_id=task_id, + summary=f"ADB Command: {cmd_line[:50]}...", + trace_id=trace_id, + ) + if stdout_data: + ctx.data_engine.stream_output(trace_id, "".join(stdout_data)) + + if ctx: + await bg_task.start(ctx) + + response_msg = f"ADB command was sent to the background as a task.\nTaskId: {task_id}\n" + if run_persistent: + response_msg += f"TerminalID: {terminal_id}\n" + response_msg += "Use the 'manage_task' tool to check its status or terminate it." + return response_msg + + +# Universal tool instance & aliases +run_adb_command = RunAdbCommandTool() +RunAdbCommand = RunAdbCommandTool +ToolRegistry.register(run_adb_command) + + +def get_run_adb_command_tool(ctx: ArtemisContext) -> BaseTool: + """Exports run_adb_command as a LangChain BaseTool.""" + return trace_langchain_tool(run_adb_command.to_langchain_tool(ctx), ctx) + + +class ManageTaskArgs(BaseModel): + """Arguments schema for managing background ADB tasks.""" + + Action: Literal["list", "list_finished", "kill", "status", "send_input"] = Field( + ..., + description=( + "The action to perform: 'list' (list running tasks)," + " 'list_finished' (list completed tasks), 'kill', 'status'," + " 'send_input'." + ), + ) + TaskId: str | None = Field( + default=None, + description=("The task ID to manage (required for kill, status, and send_input)."), + ) + Input: str | None = Field( + default=None, + description=("The input to send to the task (required when Action is 'send_input')."), + ) + + +class ManageTaskTool(ArtemisTool): + """Universal tool for managing background ADB shell tasks.""" + + def __init__(self): + super().__init__( + name="manage_task", + description=("[SHELL] Manage background ADB shell tasks launched via run_adb_command."), + args_schema=ManageTaskArgs, + category="system", + ) + + # pylint: disable=too-many-arguments,too-many-locals,too-many-branches,too-many-statements,too-many-return-statements,too-many-positional-arguments + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + ctx: ArtemisContext | None = None, + Action: Literal["list", "list_finished", "kill", "status", "send_input"] = "list", + TaskId: str | None = None, + Input: str | None = None, + **kwargs: Any, + ) -> str: + action = Action or kwargs.get("action", "list") + task_id = TaskId if TaskId is not None else kwargs.get("task_id") + input_str = Input if Input is not None else kwargs.get("input") + + if action == "list": + active_tasks = [] + for tid, t in _BACKGROUND_TASKS.items(): + active_tasks.append( + f"- {tid}: Command='{t.command[:40]}...', Phone" + f" Cwd='{t.cwd}', TerminalID='{t.terminal_id}'" + ) + if not active_tasks: + return "No active background tasks running." + return "Active background tasks:\n" + "\n".join(active_tasks) + + if action == "list_finished": + finished_tasks = [] + for tid, t in _FINISHED_TASKS_LOGS.items(): + finished_tasks.append( + f"- {tid}: Command='{t['command'][:40]}...', Status='{t['status']}'" + ) + if not finished_tasks: + return "No recently finished tasks." + return "Recently finished tasks (up to 10):\n" + "\n".join(finished_tasks) + + if not task_id: + return "Error: TaskId is required for status, kill, or send_input actions." + + task = _BACKGROUND_TASKS.get(task_id) + + if action == "status": + task_info = _get_task_info(task_id, ctx) + if not task_info: + return f"Task {task_id} not found." + + output_text = task_info.get("output", "") + if "===ENV_START===" in output_text: + output_text = output_text.split("===ENV_START===")[0] + + intro = ( + f"Task: {task_id}\n" + f"Status: {task_info.get('status')}\n" + f"Command: {task_info.get('command')}\n" + f"Cwd: {task_info.get('cwd')}\n" + f"TerminalID: {task_info.get('terminal_id')}" + ) + + if _is_output_long(output_text): + return _format_long_output_response(task_id, output_text, intro) + + return f"{intro}\nAccumulated output:\n{output_text}" + + if action == "kill": + if not task: + return f"Task {task_id} is not active or already finished." + try: + task.process.terminate() + await task.process.wait() + task.status = "killed" + if ctx and ctx.data_engine: + ctx.data_engine.unregister_background_task(task_id, "killed") + if task_id in _BACKGROUND_TASKS: + del _BACKGROUND_TASKS[task_id] + return f"Task {task_id} successfully terminated." + except Exception as e: # pylint: disable=broad-exception-caught + return f"Failed to terminate task {task_id}: {e}" + + if action == "send_input": + if not task: + return f"Task {task_id} is not active." + if not input_str: + return "Error: Input is required for send_input action." + try: + task.process.stdin.write(input_str.encode()) + await task.process.stdin.drain() + return f"Input successfully sent to task {task_id}." + except Exception as e: # pylint: disable=broad-exception-caught + return f"Failed to send input: {e}" + + return "Unsupported action." + + +# Universal tool instance & aliases +manage_task = ManageTaskTool() +ManageTask = ManageTaskTool +ToolRegistry.register(manage_task) + + +def get_manage_task_tool(ctx: ArtemisContext) -> BaseTool: + """Exports manage_task as a LangChain BaseTool.""" + return trace_langchain_tool(manage_task.to_langchain_tool(ctx), ctx) + + +class RunShortAdbCommandArgs(BaseModel): + """Arguments schema for short ADB shell commands.""" + + CommandLine: str = Field( + ..., + description=( + "The exact shell command line string to execute inside the Android device shell." + ), + ) + + +class RunShortAdbCommandTool(ArtemisTool): + """Universal tool for executing short, synchronous ADB shell commands.""" + + def __init__(self, name: str = "run_short_adb_command"): + super().__init__( + name=name, + description=( + "[SHELL] Executes a short, synchronous shell command directly on the Android " + "mobile device via ADB shell. Returns the command output. Times out after 3 " + "seconds. Do not use this tool to execute actions that change the phone state " + "such as clicking, swiping, or navigating back." + ), + args_schema=RunShortAdbCommandArgs, + category="system", + ) + + # pylint: disable=too-many-locals,too-many-return-statements + async def execute( + self, + driver: BaseDeviceDriver | None = None, + ctx: ArtemisContext | None = None, + CommandLine: str | None = None, + **kwargs: Any, + ) -> str: + cmd_line = ( + CommandLine + if CommandLine is not None + else (kwargs.get("command_line") or kwargs.get("command") or kwargs.get("cmd") or "") + ) + + # Fallback to driver's execute_shell if no full ArtemisContext is available + if ( + (ctx is None or not hasattr(ctx, "device") or not ctx.device) + and driver is not None + and hasattr(driver, "execute_shell") + ): + return await driver.execute_shell(cmd_line, timeout_seconds=3.0) + + device_id = None + if ctx and hasattr(ctx, "device") and ctx.device: + device_id = getattr(ctx.device, "device_id", None) + elif driver and hasattr(driver, "device_id"): + device_id = getattr(driver, "device_id", None) + if not device_id: + device_id = "default_device" + + run_env = os.environ.copy() + phone_script = f"cd /data/local/tmp\n{cmd_line}\n" + + adb_client = getattr(ctx, "adb_client", None) if ctx else None + if os.environ.get("ARTEMIS_CLOUD_MODE") == "1" or ( + adb_client is not None and hasattr(adb_client, "_bridge") + ): + try: + device = get_adb_device(ctx) + output = device.shell(phone_script) + return f"ADB command completed.\nOutput:\n{output}" + except Exception as e: # pylint: disable=broad-exception-caught + return f"Error running command: {e}" + + try: + process = await asyncio.create_subprocess_exec( + "adb", + "-s", + device_id, + "shell", + phone_script, + cwd=os.getcwd(), + env=run_env, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, + stdin=asyncio.subprocess.PIPE, + ) + except Exception as e: # pylint: disable=broad-exception-caught + logger.error(f"Failed to spawn adb shell process for command '{cmd_line}': {e}") + return f"Failed to spawn adb process: {e}" + + try: + stdout_data, _ = await asyncio.wait_for(process.communicate(), timeout=3.0) + exit_code = process.returncode + output_text = stdout_data.decode() + + return f"ADB command completed with exit code {exit_code}.\nOutput:\n{output_text}" + except TimeoutError: + try: + process.terminate() + await process.wait() + except Exception: # pylint: disable=broad-exception-caught + pass + return "Error: Command timed out after 3 seconds." + except Exception as e: # pylint: disable=broad-exception-caught + return f"Error running command: {e}" + + +# Universal tool instance & aliases +run_short_adb_command = RunShortAdbCommandTool() +RunShortAdbCommand = RunShortAdbCommandTool +ToolRegistry.register(run_short_adb_command) + + +def get_run_short_adb_command_tool(ctx: ArtemisContext) -> BaseTool: + """Exports run_short_adb_command as a LangChain BaseTool.""" + return trace_langchain_tool( + run_short_adb_command.to_langchain_tool(ctx, name="run_adb_command"), + ctx, + ) + + +# Expose Wrappers +run_adb_command_wrapper = ToolWrapper( + tool_fn_getter=get_run_adb_command_tool, + on_success_fn=lambda output: output, + on_failure_fn=lambda err: f"ADB command failed: {err}", +) + +manage_task_wrapper = ToolWrapper( + tool_fn_getter=get_manage_task_tool, + on_success_fn=lambda output: output, + on_failure_fn=lambda err: f"Task management failed: {err}", +) + + +class AnalyzeTaskOutputArgs(BaseModel): + """Arguments schema for output analysis requests.""" + + TaskId: str = Field( + ..., + description="The task ID whose output you want to analyze.", + ) + Query: str = Field( + ..., + description=( + "The question or query to ask the output analyzer about the task's full output log." + ), + ) + + +class AnalyzeTaskOutputTool(ArtemisTool): + """Universal tool for analyzing the full, long output of an ADB task.""" + + def __init__(self): + super().__init__( + name="analyze_task_output", + description=( + "[ANALYSIS] Analyzes the full, long output of a completed or running ADB task. " + "Use this tool when you receive a message indicating that a task's output " + "is too long and was truncated." + ), + args_schema=AnalyzeTaskOutputArgs, + category="system", + ) + + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + ctx: ArtemisContext | None = None, + TaskId: str | None = None, + Query: str | None = None, + **kwargs: Any, + ) -> str: + task_id = TaskId or kwargs.get("task_id") or "" + query = Query or kwargs.get("query") or "" + + task_info = _get_task_info(task_id, ctx) + if not task_info: + return f"Error: Task {task_id} not found." + + analyzer = TaskOutputAnalyzerNode(ctx) + result = await analyzer.run( + command=task_info.get("command", ""), + output_text=task_info.get("output", ""), + query=query, + ) + return f"Analysis result for Task {task_id}:\n{result}" + + +# Universal tool instance & aliases +analyze_task_output = AnalyzeTaskOutputTool() +AnalyzeTaskOutput = AnalyzeTaskOutputTool +ToolRegistry.register(analyze_task_output) + + +def get_analyze_task_output_tool(ctx: ArtemisContext) -> BaseTool: + """Exports analyze_task_output as a LangChain BaseTool.""" + return trace_langchain_tool(analyze_task_output.to_langchain_tool(ctx), ctx) + + +analyze_task_output_wrapper = ToolWrapper( + tool_fn_getter=get_analyze_task_output_tool, + on_success_fn=lambda output: output, + on_failure_fn=lambda err: f"Task output analysis failed: {err}", +) diff --git a/artemis/tools/committee_prompts.json b/artemis/tools/committee_prompts.json new file mode 100644 index 0000000..253df2f --- /dev/null +++ b/artemis/tools/committee_prompts.json @@ -0,0 +1,5 @@ +{ + "planner_avatar": "You are the **Planner Avatar** (Meeting Host) in a committee resolving a task plan failure. Your role is to lead the meeting and advocate for a robust and correct high-level plan. You are aware of:\n- **History Analyzer**: The trajectory specialist who can analyze execution history and read notes. They have tools like `get_step_details` and `read_note`.\n- **Diagnoser**: The forensics specialist who can analyze logs and videos to find why things failed. They have tools like `video_analyzer` and `analyze_logs`.\nYour master (Master Planner) has given you a specific directive for this meeting, which is at the top of the blackboard.\n**Your Task**:\n1. In the first round, you must open the meeting, state the problem based on the master's directive, and call on the other experts (Diagnoser or History Analyzer) to provide their analysis.\n2. In subsequent rounds, evaluate their suggestions and push for a concrete, revised plan of subgoals.\n3. In the final round, summarize the discussion and output the final recommended plan revisions and actions.", + "history_analyzer_expert": "You are the **History Analyzer** (Trajectory Specialist) in a committee resolving a task failure. Your role is to analyze the execution history of the current session, understand what has been tried, what succeeded, and what failed, and check persistent notes. You are aware of:\n- **Planner Avatar**: The meeting host who needs a concrete plan.\n- **Diagnoser**: The forensics specialist who analyzes logs and videos.\nYou have tools to retrieve full-granularity step details and read notes. Wait for the Planner Avatar to call on you or respond to the current state of the blackboard. Be concise.", + "diagnoser_expert": "You are the **Diagnoser** (Forensics Specialist) in a committee resolving a UI automation failure. Your role is to analyze what went wrong in the past, look for error patterns in logs, and understand system state transitions. You are aware of:\n- **Planner Avatar**: The meeting host who needs a concrete plan.\n- **History Analyzer**: The trajectory specialist who analyzes history and notes.\nYou have tools to analyze video and logs. Wait for the Planner Avatar to call on you or respond to the current state of the blackboard. Be concise." +} diff --git a/artemis/tools/committee_tool.py b/artemis/tools/committee_tool.py new file mode 100644 index 0000000..2d55068 --- /dev/null +++ b/artemis/tools/committee_tool.py @@ -0,0 +1,427 @@ +# pylint: disable=too-many-statements +# 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. + +"""Committee tool for multi-agent deliberation and failure recovery.""" + +import asyncio +import base64 +import json +from pathlib import Path +from typing import Any +from uuid import uuid4 + +from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage +from langchain_core.tools import BaseTool +from pydantic import BaseModel, Field + +from artemis.agents.history_analyzer.history_analyzer import HistoryAnalyzer +from artemis.context import ArtemisContext +from artemis.data_engine.trace import ( + CURRENT_TRACE_ID, + TraceSpan, + trace_langchain_tool, +) +from artemis.drivers.base import BaseDeviceDriver +from artemis.graph.state import State +from artemis.services.llm import get_llm, invoke_llm_with_timeout_message +from artemis.tools.base import ArtemisTool, ToolCategory, ToolRegistry +from artemis.tools.log_tool import get_analyze_logs_tool +from artemis.tools.scratchpad import get_list_notes_tool, get_read_note_tool +from artemis.tools.tool_wrapper import ( + ToolWrapper, + get_tool_result_content, + invoke_tool_with_injection, +) +from artemis.tools.types import CyFunctionDetector +from artemis.tools.video_tool import get_video_analyzer_tool +from artemis.utils.logger import get_logger +from artemis.utils.notes import get_note_file_path, get_notes_dir +from artemis.utils.task_tree import ( + build_plan_and_history, + get_active_subgoal_hashes, + get_recent_subgoal_hashes, +) + +logger = get_logger(__name__) + + +class AskCommitteeArgs(BaseModel): + """Arguments schema for invoking the committee tool.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + avatar_directive: str = Field( + ..., + description="Directive to frame and steer the debate for your avatar.", + ) + + +ASK_COMMITTEE_DOCSTRING = ( + "[COMMITTEE] Call this tool to summon a committee of sub-agents to debate" + " and resolve complex failures.\n\n" + "Your avatar will host the meeting and drive it based on your directive.\n" + "The result will be a synthesized recommendation." +) + + +class AskCommitteeTool(ArtemisTool): + """Universal tool for summoning a committee of sub-agents to debate and resolve failures.""" + + def __init__(self, category: ToolCategory = "custom"): + super().__init__( + name="ask_committee", + description=ASK_COMMITTEE_DOCSTRING, + args_schema=AskCommitteeArgs, + category=category, + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + ctx: ArtemisContext | None = None, + avatar_directive: str = "", + tool_call_id: str | None = None, # pylint: disable=unused-argument + state: State | None = None, + **kwargs: Any, + ) -> str: + directive = avatar_directive or kwargs.get("AvatarDirective") or "" + return await _run_committee_logic( + ctx=ctx, + state=state, + avatar_directive=directive, + ) + + +# Universal tool instance & aliases +ask_committee = AskCommitteeTool() +AskCommittee = AskCommitteeTool +AskCommitteeToolAlias = AskCommitteeTool +ToolRegistry.register(ask_committee) + + +def get_ask_committee_tool(ctx: ArtemisContext) -> BaseTool: + """Exports ask_committee as a LangChain BaseTool.""" + return trace_langchain_tool(ask_committee.to_langchain_tool(ctx), ctx) + + +# pylint: disable=too-many-branches,too-many-locals,too-many-statements +async def _execute_committee( + ctx: ArtemisContext, + state: State | None, + avatar_directive: str, +) -> str: + """Executes committee debate loop and returns the final conclusion.""" + with TraceSpan(name="committee", trace_type="agent", ctx=ctx): + logger.info(f"Committee invoked with directive: {avatar_directive}") + + # 1. Access or create blackboard file + if ctx.data_engine: + notes_dir = get_notes_dir(ctx.data_engine.base_dir) + else: + notes_dir = Path(".") / "notes" + notes_dir.mkdir(parents=True, exist_ok=True) + + # Load latest screenshot as base64 to provide multimodal vision to the committee + screenshot_b64 = None + screenshot_path = getattr(state, "latest_screenshot", None) if state is not None else None + if screenshot_path and Path(screenshot_path).exists(): + try: + with open(screenshot_path, "rb") as f: + screenshot_b64 = base64.b64encode(f.read()).decode("utf-8") + logger.info(f"Committee successfully loaded screenshot from: {screenshot_path}") + except Exception as e: # pylint: disable=broad-exception-caught + logger.error(f"Failed to read screenshot in committee: {e}") + + trace_id = CURRENT_TRACE_ID.get() or str(uuid4()) + blackboard_path = notes_dir / f"blackboard_{trace_id}.md" + + # Initialize blackboard with goal and directive + initial_goal = ( + getattr(state, "initial_goal", "Not specified.") + if state is not None + else "Not specified." + ) + blackboard_content = ( + "# Committee Discussion Blackboard\n\n" + f"## User's Initial Goal\n{initial_goal}\n\n" + f"## Master's Directive\n{avatar_directive}\n\n" + ) + blackboard_path.write_text(blackboard_content, encoding="utf-8") + + # 2. Extract history + history_summary = "## Previous Context\n" + if ctx.data_engine: + steps = ctx.data_engine.get_agent_friendly_steps() + if steps: + active_subgoal_hash = "default" + if ctx.data_engine: + task_plan_path = get_note_file_path(ctx.data_engine.base_dir, "task_plan") + else: + task_plan_path = notes_dir / "task_plan.md" + task_plan = "No task plan yet." + keep_hashes = None + if task_plan_path.exists(): + try: + task_plan = task_plan_path.read_text(encoding="utf-8") + + active_subgoal_hash, _ = get_active_subgoal_hashes(task_plan) + keep_hashes = get_recent_subgoal_hashes( + steps, + active_subgoal_hash, + ctx.data_engine.base_dir, + ) + except Exception as e: # pylint: disable=broad-exception-caught + logger.error(f"Failed to parse active subgoal in committee: {e}") + + try: + plan_and_history = build_plan_and_history( + task_plan, + steps, + active_subgoal_hash, + keep_subgoal_hashes=keep_hashes, + last_n_detailed=3, + ) + history_summary += plan_and_history + except Exception as e: # pylint: disable=broad-exception-caught + logger.error(f"Failed to build plan and history in committee: {e}") + history_summary += f"Error building plan and history: {e}" + + # Append immediate failure reason if present in state + replan_reason = ( + getattr(state, "operator_replan_reason", None) if state is not None else None + ) + if replan_reason: + history_summary += f"\n\n### Current Failure Reason\n{replan_reason}\n" + + # Read failed plans history to ensure the committee avoids + # repeating failed strategies + failed_plans_history = "" + failed_history_path = notes_dir / "failed_plans_history.md" + if failed_history_path.exists(): + try: + failed_plans_history = failed_history_path.read_text(encoding="utf-8") + except Exception as e: # pylint: disable=broad-exception-caught + logger.error(f"Failed to read failed plans history in committee: {e}") + if failed_plans_history: + history_summary += ( + f"\n\n### Previous Failed Plan Attempts\n{failed_plans_history}\n" + ) + + with open(blackboard_path, "a", encoding="utf-8") as f: + f.write(history_summary + "\n\n") + + # 3. Load prompts from JSON + prompts = {} + prompts_path = Path(__file__).parent / "committee_prompts.json" + if prompts_path.exists(): + try: + prompts = json.loads(prompts_path.read_text(encoding="utf-8")) + except Exception as e: # pylint: disable=broad-exception-caught + logger.error(f"Failed to load committee prompts: {e}") + + planner_prompt = prompts.get("planner_avatar", "You are the Planner Avatar.") + history_analyzer_prompt = prompts.get( + "history_analyzer_expert", "You are the History Analyzer." + ) + diagnoser_prompt = prompts.get("diagnoser_expert", "You are the Diagnoser Expert.") + + # 4. Prepare Specialized Tools for Members + + # History Analyzer Tools + analyzer_inst = HistoryAnalyzer(ctx) + history_steps = ctx.data_engine.get_agent_friendly_steps() if ctx.data_engine else [] + # pylint: disable=protected-access + details_tool = analyzer_inst._get_step_details_tool(history_steps) + read_note_tool_hist = get_read_note_tool(ctx) + list_notes_tool_hist = get_list_notes_tool(ctx) + + history_analyzer_tools = [ + trace_langchain_tool(details_tool, ctx), + trace_langchain_tool(read_note_tool_hist, ctx), + trace_langchain_tool(list_notes_tool_hist, ctx), + ] + + # Diagnoser Tools + video_tool = get_video_analyzer_tool(ctx, role="diagnoser") + log_tool = get_analyze_logs_tool(ctx) + read_note_tool = get_read_note_tool(ctx) + list_notes_tool = get_list_notes_tool(ctx) + + diagnoser_tools = [ + trace_langchain_tool(video_tool, ctx), + trace_langchain_tool(log_tool, ctx), + trace_langchain_tool(read_note_tool, ctx), + trace_langchain_tool(list_notes_tool, ctx), + ] + + # 5. Instantiate LLMs + llm_pl = get_llm(ctx=ctx, name="planner_avatar") + + llm_hist = get_llm(ctx=ctx, name="history_analyzer_expert") + llm_hist = llm_hist.bind_tools(tools=history_analyzer_tools) + + llm_diag = get_llm(ctx=ctx, name="diagnoser_expert") + llm_diag = llm_diag.bind_tools(tools=diagnoser_tools) + + # 6. Debate Loop (configured rounds) + exec_setup = getattr(ctx, "execution_setup", None) + rounds = getattr(exec_setup, "committee_debate_rounds", 2) if exec_setup else 2 + + async def run_agent_turn(llm, system_prompt, tools, agent_name): + current_blackboard = blackboard_path.read_text(encoding="utf-8") + + human_content = [ + { + "type": "text", + "text": f"Current Blackboard:\n{current_blackboard}", + } + ] + if screenshot_b64: + human_content.append( + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{screenshot_b64}"}, + } + ) + + messages = [ + SystemMessage(content=system_prompt), + HumanMessage(content=human_content), + ] + + max_iterations = 3 + for _ in range(max_iterations): + response = await invoke_llm_with_timeout_message(llm.ainvoke(messages)) + + if not response.tool_calls: + return response.content + + messages.append(response) + + async def run_tool(tc): + tool_name = tc["name"] + logger.info(f"{agent_name} requested tool: {tool_name}") + tool_to_run = next((t for t in tools if t.name == tool_name), None) + if tool_to_run: + try: + args = dict(tc["args"]) + result = await invoke_tool_with_injection( + tool=tool_to_run, + args=args, + tool_call_id=tc["id"], + state=state, + ) + return get_tool_result_content(result) + except Exception as e: # pylint: disable=broad-exception-caught + return f"Error: {e}" + else: + return f"Error: Tool {tool_name} not found" + + active_tool_calls = [ + tc for tc in response.tool_calls if tc["name"] != "google_search" + ] + tool_outputs = await asyncio.gather(*(run_tool(tc) for tc in active_tool_calls)) + + for tc, result in zip(active_tool_calls, tool_outputs): + messages.append( + ToolMessage( + tool_call_id=tc["id"], + content=result, + status="success" if not result.startswith("Error") else "error", + ) + ) + return "Error: Reached max iterations in turn." + + for r in range(1, rounds + 1): + logger.info(f"Committee Round {r}") + + # 1. Planner Avatar speaks first (Host) + pl_response = await run_agent_turn(llm_pl, planner_prompt, [], "Planner Avatar") + with open(blackboard_path, "a", encoding="utf-8") as f: + f.write(f"### Planner Avatar (Round {r})\n{pl_response}\n\n") + + # 2. Diagnoser speaks (Can use tools) + diag_response = await run_agent_turn( + llm_diag, diagnoser_prompt, diagnoser_tools, "Diagnoser" + ) + with open(blackboard_path, "a", encoding="utf-8") as f: + f.write(f"### Diagnoser (Round {r})\n{diag_response}\n\n") + + # 3. History Analyzer speaks (Can use tools) + hist_response = await run_agent_turn( + llm_hist, + history_analyzer_prompt, + history_analyzer_tools, + "History Analyzer", + ) + with open(blackboard_path, "a", encoding="utf-8") as f: + f.write(f"### History Analyzer (Round {r})\n{hist_response}\n\n") + + # Final Synthesis by Planner Avatar + current_blackboard = blackboard_path.read_text(encoding="utf-8") + + final_human_content = [ + { + "type": "text", + "text": f"Current Blackboard:\n{current_blackboard}", + } + ] + if screenshot_b64: + final_human_content.append( + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{screenshot_b64}"}, + } + ) + + msg_final = [ + SystemMessage( + content=( + "You are the Planner Avatar. Summarize the discussion" + " and output the final recommended plan revisions and" + " action plan based on your master's directive." + ) + ), + HumanMessage(content=final_human_content), + ] + resp_final = await invoke_llm_with_timeout_message(llm_pl.ainvoke(msg_final)) + + with open(blackboard_path, "a", encoding="utf-8") as f: + f.write(f"## Final Conclusion\n{resp_final.content}\n") + + return resp_final.content + + +async def _run_committee_logic( + ctx: ArtemisContext | None, + state: State | None, + avatar_directive: str, +) -> str: + """Executes multi-agent committee debate logic.""" + if ctx is None: + return "Error: ArtemisContext is required for ask_committee." + + try: + return await asyncio.wait_for(_execute_committee(ctx, state, avatar_directive), timeout=300) + except TimeoutError: + logger.error("Committee timed out after 300 seconds.") + return "Error: Committee timed out after 300 seconds." + + +ask_committee_wrapper = ToolWrapper( + tool_fn_getter=get_ask_committee_tool, + on_success_fn=lambda output: f"Committee concluded:\n{output}", + on_failure_fn=lambda error: f"Committee failed: {error}", +) diff --git a/artemis/tools/diagnoser_submit_answer_tool.py b/artemis/tools/diagnoser_submit_answer_tool.py new file mode 100644 index 0000000..ee79f69 --- /dev/null +++ b/artemis/tools/diagnoser_submit_answer_tool.py @@ -0,0 +1,88 @@ +# 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. + +"""Universal tool for submitting final diagnosis and actionable recovery advice.""" + +from typing import Any + +from langchain_core.tools import BaseTool +from pydantic import BaseModel, Field + +from artemis.context import ArtemisContext +from artemis.data_engine.trace import trace_langchain_tool +from artemis.drivers.base import BaseDeviceDriver +from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +# pylint: disable=too-few-public-methods +class SubmitAnswerArgs(BaseModel): + """Arguments schema for submitting final diagnoser answer.""" + + analysis: str = Field( + ..., + description=( + "A detailed and concise explanation of the root cause of the" + " failure or query. Answer in one full paragraph." + ), + ) + actionable_steps: list[str] = Field( + default_factory=list, + description=( + "A list of concrete recovery instructions or next steps. Each" + " step is a short sentence. If no actionable steps are" + " available, return an empty list." + ), + ) + + +class SubmitAnswerTool(ArtemisTool): + """Universal tool for submitting final diagnosis and actionable recovery advice.""" + + def __init__(self): + super().__init__( + name="submit_answer", + description=( + "[Finalization] Submits the final diagnosis and actionable recovery advice. " + "Call this tool when you have pinpointed the root cause." + ), + args_schema=SubmitAnswerArgs, + category="custom", + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + ctx: ArtemisContext | None = None, # pylint: disable=unused-argument + analysis: str | None = None, # pylint: disable=unused-argument + actionable_steps: list[str] | None = None, # pylint: disable=unused-argument + **kwargs: Any, + ) -> str: + logger.info("Diagnoser submitted answer via tool.") + return "Answer submitted successfully." + + +# Universal tool instance & aliases +submit_answer = SubmitAnswerTool() +SubmitAnswer = SubmitAnswerTool +DiagnoserSubmitAnswerTool = SubmitAnswerTool +ToolRegistry.register(submit_answer) + + +def get_submit_answer_tool(ctx: ArtemisContext) -> BaseTool: + """Exports submit_answer as a LangChain BaseTool.""" + return trace_langchain_tool(submit_answer.to_langchain_tool(ctx), ctx) diff --git a/artemis/tools/diagnostic_tool.py b/artemis/tools/diagnostic_tool.py new file mode 100644 index 0000000..524019a --- /dev/null +++ b/artemis/tools/diagnostic_tool.py @@ -0,0 +1,120 @@ +# 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. + +"""Universal diagnostic agent tool for deep system logs and video inspection.""" + +from typing import Any + +from langchain_core.tools import BaseTool +from pydantic import BaseModel, Field + +from artemis.agents.diagnoser.diagnoser import Diagnoser +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, ToolRegistry +from artemis.tools.tool_wrapper import ToolWrapper +from artemis.tools.types import CyFunctionDetector +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +# pylint: disable=too-few-public-methods +class AskDiagnoserArgs(BaseModel): + """Arguments schema for invoking the diagnostic subagent.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + query: str = Field( + ..., + description=( + "The diagnostic query or what you want to investigate. Please" + " provide a specific, narrow time range if applicable." + ), + ) + + +async def _run_diagnoser_logic( + ctx: ArtemisContext | None, + state: State | None, + query: str, +) -> str: + logger.info(f"ask_diagnoser tool called with query: {query}") + if ctx is None: + return "Error: ArtemisContext is required for Diagnoser." + agent = Diagnoser(ctx) + try: + result = await agent.run(query, state) + return f"Diagnoser replied:\n{result}" + except Exception as e: # pylint: disable=broad-exception-caught + logger.error(f"Failed to run diagnoser: {e}") + return f"Diagnoser failed: {e}" + + +class AskDiagnoserTool(ArtemisTool): + """Universal tool for invoking the Diagnoser subagent to inspect logs and videos.""" + + def __init__(self): + super().__init__( + name="ask_diagnoser", + description=( + "[DIAGNOSTIC] Analyzes system logs (ADB) and video recordings to " + "pinpoint the root cause of any step failure, screen freeze, or " + "unexpected UI drift, and directly executes adb commands. Call this " + "tool proactively whenever an action fails, status stagnates, or you " + "are unsure why progress is blocked, to diagnose the issue before " + "committing to physical actions. Note: When querying video files, " + "specify a narrow timeframe (e.g., '[10s, 20s]') for faster analysis." + ), + args_schema=AskDiagnoserArgs, + category="custom", + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + ctx: ArtemisContext | None = None, + query: str | None = None, + state: State | None = None, + **kwargs: Any, + ) -> str: + q = query or kwargs.get("Query") or "" + st = state if state is not None else kwargs.get("state") + return await _run_diagnoser_logic(ctx, st, q) + + +# Universal tool instance & aliases +ask_diagnoser = AskDiagnoserTool() +AskDiagnoser = AskDiagnoserTool +DiagnosticTool = AskDiagnoserTool +ToolRegistry.register(ask_diagnoser) + + +def get_ask_diagnostic_agent_tool(ctx: ArtemisContext) -> BaseTool: + """Exports ask_diagnoser as a LangChain BaseTool.""" + return trace_langchain_tool(ask_diagnoser.to_langchain_tool(ctx), ctx) + + +def get_ask_diagnoser_tool(ctx: ArtemisContext) -> BaseTool: + """Exports ask_diagnoser as a LangChain BaseTool (standard naming alias).""" + return trace_langchain_tool(ask_diagnoser.to_langchain_tool(ctx), ctx) + + +ask_diagnoser_wrapper = ToolWrapper( + tool_fn_getter=get_ask_diagnostic_agent_tool, + on_success_fn=lambda output: f"Diagnoser replied:\n{output}", + on_failure_fn=lambda error: f"Diagnoser failed: {error}", +) diff --git a/artemis/tools/explorer_tool.py b/artemis/tools/explorer_tool.py new file mode 100644 index 0000000..a696e5b --- /dev/null +++ b/artemis/tools/explorer_tool.py @@ -0,0 +1,404 @@ +# 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 glob +import json +import os +from pathlib import Path +import re +from typing import Any, Literal + +from langchain_core.tools import BaseTool +from pydantic import BaseModel, Field, create_model + +from artemis.agents.explorer.constants import EXPLORE_DESCRIPTIONS +from artemis.agents.explorer.explorer import Explorer +from artemis.agents.object_detector.object_detector import _run_object_detection +from artemis.config import ExplorerVersion, resolve_explorer_version, settings +from artemis.context import ArtemisContext +from artemis.data_engine.trace import trace, 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.tools.types import CyFunctionDetector +from artemis.utils.logger import get_logger +from artemis.utils.visualization import draw_dots + + +class AskExplorerArgs(BaseModel): + """Arguments schema for invoking the explorer subagent.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + query: str = Field( + ..., + description=("The target element or information to search for, including descriptions."), + ) + context_feedback: str = Field( + "", + description=("Feedback from previous failed attempts or specific instructions"), + ) + version: Literal["flash", "pro", "ultra"] | None = Field( + default=None, + description="Explicit explorer version to use (flash, pro, or ultra).", + ) + + +logger = get_logger(__name__) + + +class AskExplorerTool(ArtemisTool): + """Universal tool for visually parsing and locating UI elements.""" + + def __init__( + self, + version: ExplorerVersion | None = None, + agent_name: str = "operator", + description: str | None = None, + category: ToolCategory = "explorer", + ): + self.version = version + self.agent_name = agent_name + resolved_v = version or "pro" + explore_info = EXPLORE_DESCRIPTIONS.get(resolved_v, EXPLORE_DESCRIPTIONS["pro"]) + desc = description or explore_info["description"] + if "{max_iterations}" in desc: + max_iterations = 8 if resolved_v == "ultra" else 3 + desc = desc.format(max_iterations=max_iterations) + + super().__init__( + name="ask_explorer", + description=desc, + args_schema=AskExplorerArgs, + category=category, + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + ctx: ArtemisContext | None = None, + query: str = "", + context_feedback: str = "", + version: Literal["flash", "pro", "ultra"] | None = None, + tool_call_id: str | None = None, # pylint: disable=unused-argument + state: State | None = None, + **kwargs: Any, + ) -> Any: + q = query or kwargs.get("Query") or "" + cf = context_feedback or kwargs.get("ContextFeedback") or "" + v = version or self.version + resolved_v = ( + resolve_explorer_version(ctx, explicit_version=v, agent_or_profile_name=self.agent_name) + if ctx + else (v or "pro") + ) + return await _run_explorer_logic( + ctx=ctx, + state=state, + query=q, + context_feedback=cf, + version=resolved_v, + ) + + +# Universal tool instance & aliases +ask_explorer = AskExplorerTool() +AskExplorer = AskExplorerTool +AskExplorerToolAlias = AskExplorerTool +ToolRegistry.register(ask_explorer) + + +def get_ask_explorer_tool( + ctx: ArtemisContext, + version: ExplorerVersion | None = None, + agent_name: str = "operator", +) -> BaseTool: + """Exports ask_explorer as a LangChain BaseTool with dynamic schema.""" + resolved_version = resolve_explorer_version( + ctx, explicit_version=version, agent_or_profile_name=agent_name + ) + + explore_info = EXPLORE_DESCRIPTIONS.get(resolved_version, EXPLORE_DESCRIPTIONS["pro"]) + description = explore_info["description"] + if "{max_iterations}" in description: + max_iterations = 8 if resolved_version == "ultra" else 3 + description = description.format(max_iterations=max_iterations) + query_description = explore_info["query_description"] + + # Dynamically build argument schema so Operator does not see + # context_feedback when flash is active + fields = { + "query": (str, Field(..., description=query_description)), + } + if resolved_version != "flash": + fields["context_feedback"] = ( + str, + Field( + "", + description=("Feedback from previous failed attempts or specific instructions"), + ), + ) + + # pylint: disable=invalid-name + DynamicAskExplorerArgs = create_model("AskExplorerArgs", **fields) + + tool_instance = AskExplorerTool( + version=resolved_version, + agent_name=agent_name, + description=description, + ) + tool_instance.args_schema = DynamicAskExplorerArgs + + return trace_langchain_tool(tool_instance.to_langchain_tool(ctx), ctx) + + +# pylint: disable=too-many-branches,too-many-locals,too-many-statements +async def _run_explorer_logic( + ctx: ArtemisContext, + state: State, + query: str, + context_feedback: str, + version: Literal["flash", "pro", "ultra"], +) -> Any: + if version == "flash": + logger.info( + "Explorer [version='flash'] intercepted -> Direct one-shot object" + " detection (context_feedback ignored)." + ) + screenshot_path = state.latest_screenshot + if not screenshot_path: + return "Explorer (Flash) error: No latest_screenshot available in state." + + queries = [q.strip() for q in query.split("|") if q.strip()] + if not queries: + queries = [query.strip()] + + try: + prompt_path = Path(__file__).parent.parent.joinpath( + "agents", "object_detector", "object_detector.json" + ) + with open(prompt_path, encoding="utf-8") as f: + cfg = json.load(f) + templates = [f"{t}\n\n{cfg.get('instructions', '')}" for t in cfg.get("templates", [])] + + result_dict = await _run_object_detection( + ctx, screenshot_path, queries, templates, global_timeout=15.0 + ) + detected_items = result_dict.get("detected", []) + candidates = [] + for item in detected_items: + if isinstance(item, dict) and "point" in item and len(item["point"]) == 2: + nx, ny = item["point"] + lbl = item.get("label", query) + candidates.append( + { + "label": lbl, + "coords": [int(nx), int(ny)], + "description": f"{lbl} (flash detection)", + } + ) + + if not candidates: + failed_queries = result_dict.get("failed", queries) + explorer_output = json.dumps( + { + "candidates": [], + "fallback_message": ( + "Flash one-shot detection could not locate:" + f" {', '.join(failed_queries)}. Consider checking" + " visual visibility or layout." + ), + } + ) + else: + explorer_output = json.dumps({"candidates": candidates}) + + except Exception as e: # pylint: disable=broad-exception-caught + logger.error(f"Flash one-shot routing failed: {e}") + explorer_output = json.dumps( + { + "candidates": [], + "fallback_message": f"Flash detection error: {e}", + } + ) + else: + + @trace(type="agent", name="explorer", ctx=ctx) + async def run_explorer_agent(q: str, cf: str, screenshot_path: str, state: State) -> str: + agent = Explorer(ctx) + env_cache = os.getenv("ARTEMIS_EXPLORER_CACHING", "").lower() + if env_cache in ["true", "false"]: + enable_caching = env_cache == "true" + else: + agent_cfg = getattr(ctx, "agent_config", None) + if ( + agent_cfg + and hasattr(agent_cfg, "explorer") + and hasattr(agent_cfg.explorer, "caching") + ): + enable_caching = agent_cfg.explorer.caching + elif ( + ctx and ctx.execution_setup and hasattr(ctx.execution_setup, "explorer_caching") + ): + enable_caching = ctx.execution_setup.explorer_caching + else: + enable_caching = getattr(settings, "EXPLORER_CACHING", True) + return await agent.run( + q, + cf, + screenshot_path, + state, + enable_caching=enable_caching, + version=version, + ) + + screenshot_path = state.latest_screenshot + explorer_output = await run_explorer_agent(query, context_feedback, screenshot_path, state) + + try: + data = json.loads(explorer_output) + candidates = data.get("candidates", []) + fallback_message = data.get("fallback_message", "") + + if candidates: + # 1. Read Operator's active coordinates list + indexed_points = getattr(state, "indexed_points", None) + if indexed_points is None: + indexed_points = [] + state.indexed_points = indexed_points + + indexed_elements = getattr(state, "indexed_elements", None) + if indexed_elements is None: + indexed_elements = [] + state.indexed_elements = indexed_elements + + width = 1080 + 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): + width = w_raw + height = h_raw + else: + if ctx.device and getattr(ctx.device, "device_width", None): + width = ctx.device.device_width + if ctx.device and getattr(ctx.device, "device_height", None): + height = ctx.device.device_height + + registered_info = [] + for cand in candidates: + coords = cand.get("coords") + desc = cand.get("description", "Explorer candidate") + if coords and len(coords) == 2: + nx, ny = coords + pixel_x = int(max(0, min(width, nx * width / 1000))) + pixel_y = int(max(0, min(height, ny * height / 1000))) + + indexed_points.append([pixel_x, pixel_y]) + new_idx = len(indexed_points) + + # Update indexed_elements in sync + indexed_elements.append( + { + "index": new_idx, + "center": [pixel_x, pixel_y], + "text": desc, + "bounds": None, + "class": "ExplorerCandidate", + "resource_id": None, + "is_ocr": False, + } + ) + + registered_info.append(f"- [{new_idx}] '{desc}' at coordinate {coords}") + + response_msg = ( + "Explorer successfully located the following candidate(s):\n" + + "\n".join(registered_info) + ) + response_msg += ( + "\nYou can click/act on them directly by calling perform_action" + " with their respective index." + ) + + if fallback_message: + response_msg += f"\n\nAdditional Notes from Explorer: {fallback_message}" + + # 2. Generate annotated image with ALL points + + base_dir = ( + Path(ctx.data_engine.base_dir) + if ctx.data_engine and getattr(ctx.data_engine, "base_dir", None) + else settings.TRACES_PATH + ) + images_dir = base_dir / "images" + explorer_dir = images_dir / "explorer_tool" + explorer_dir.mkdir(parents=True, exist_ok=True) + + image_name = "explorer_output" + existing_files = glob.glob(str(explorer_dir / f"{image_name}_*.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 + output_path = explorer_dir / f"{image_name}_{seq}.jpg" + + all_labels = [str(i) for i in range(1, len(indexed_points) + 1)] + + try: + draw_dots( + screenshot_path, + indexed_points, + all_labels, + str(output_path), + ) + with open(output_path, "rb") as im_f: + img_bytes = im_f.read() + img_b64 = base64.b64encode(img_bytes).decode("utf-8") + + return [ + { + "type": "text", + "text": f"Explorer replied:\n{response_msg}", + }, + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{img_b64}"}, + }, + ] + except Exception as e: # pylint: disable=broad-exception-caught + logger.error(f"Failed to draw dots or read image in explorer_tool: {e}") + # Fallback to pure text if drawing fails + return f"Explorer replied:\n{response_msg}" + + elif fallback_message: + return f"Explorer could not locate the element. Message: {fallback_message}" + + except Exception as e: # pylint: disable=broad-exception-caught + logger.error(f"Failed to parse explorer output: {e}") + return f"Explorer failed to provide a valid format. Raw output: {explorer_output}" + + return f"Explorer replied:\n{explorer_output}" + + +ask_explorer_wrapper = ToolWrapper( + tool_fn_getter=get_ask_explorer_tool, + on_success_fn=lambda output: f"Explorer replied:\n{output}", + on_failure_fn=lambda error: f"Explorer failed: {error}", +) diff --git a/artemis/tools/image_processor_tool.py b/artemis/tools/image_processor_tool.py new file mode 100644 index 0000000..2df9ef4 --- /dev/null +++ b/artemis/tools/image_processor_tool.py @@ -0,0 +1,125 @@ +# 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. + +"""Universal Image Processor tool for executing image transformations.""" + +from typing import Any + +from langchain_core.tools import BaseTool +from pydantic import BaseModel, Field + +from artemis.context import ArtemisContext +from artemis.data_engine.trace import trace, trace_langchain_tool +from artemis.drivers.base import BaseDeviceDriver +from artemis.graph.state import State +from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.tool_wrapper import ToolWrapper +from artemis.tools.types import CyFunctionDetector +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +# pylint: disable=too-few-public-methods +class AskVisionCoderArgs(BaseModel): + """Arguments schema for invoking the Image Processor subagent.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + instruction: str = Field( + ..., + description=( + "Step-by-step instructions on what image processing or" + " modifications should be performed." + ), + ) + target_image_id: str = Field( + ..., + description=( + "The ID of the image in the Image Pool to start from (e.g. 'img_0', 'img_1')." + ), + ) + + +# Alias for naming consistency +AskImageProcessorArgs = AskVisionCoderArgs + + +class AskImageProcessorTool(ArtemisTool): + """Universal tool for executing image transformations via ImageProcessor subagent.""" + + def __init__(self): + super().__init__( + name="ask_image_processor", + description=( + "Call this tool to ask the image processor agent to run" + " complex image transformations or modifications." + ), + args_schema=AskVisionCoderArgs, + category="custom", + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + ctx: ArtemisContext | None = None, + instruction: str | None = None, + target_image_id: str | None = None, + state: State | None = None, + **kwargs: Any, + ) -> Any: + inst = instruction or kwargs.get("Instruction") or "" + tid = target_image_id or kwargs.get("TargetImageId") or kwargs.get("target_image_id") or "" + st = state if state is not None else kwargs.get("state") + return await _run_image_processor_logic(ctx, st, inst, tid) + + +# Universal tool instance & aliases +ask_image_processor = AskImageProcessorTool() +AskImageProcessor = AskImageProcessorTool +AskVisionCoderTool = AskImageProcessorTool +ImageProcessorTool = AskImageProcessorTool +ToolRegistry.register(ask_image_processor) + + +def get_ask_image_processor_tool(ctx: ArtemisContext) -> BaseTool: + """Exports ask_image_processor as a LangChain BaseTool.""" + return trace_langchain_tool(ask_image_processor.to_langchain_tool(ctx), ctx) + + +async def _run_image_processor_logic( + ctx: ArtemisContext | None, + state: State | None, + instruction: str, + target_image_id: str, # pylint: disable=unused-argument +) -> Any: + @trace(type="agent", name="image_processor", ctx=ctx) + async def run_image_processor_agent(inst: str, img_path: str) -> dict: + from artemis.agents.image_processor.image_processor import ( # pylint: disable=import-outside-toplevel + ImageProcessor, + ) + + agent = ImageProcessor(ctx) + return await agent.run(inst, img_path) + + screenshot_path = state.latest_screenshot if state else None + result = await run_image_processor_agent(instruction, screenshot_path) + return result + + +ask_image_processor_wrapper = ToolWrapper( + tool_fn_getter=get_ask_image_processor_tool, + on_success_fn=lambda output: f"Image Processor replied:\n{output}", + on_failure_fn=lambda error: f"Image Processor failed: {error}", +) diff --git a/artemis/tools/index.py b/artemis/tools/index.py new file mode 100644 index 0000000..80914bb --- /dev/null +++ b/artemis/tools/index.py @@ -0,0 +1,52 @@ +# 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. + +from langchain_core.tools import BaseTool + +from artemis.context import ArtemisContext +from artemis.data_engine.trace import trace_langchain_tool +from artemis.tools.tool_wrapper import CompositeToolWrapper, ToolWrapper + + +def get_tools_from_wrappers( + ctx: "ArtemisContext", + wrappers: list[ToolWrapper], +) -> list[BaseTool]: + """Instantiate and wrap LangChain tools from a list of ToolWrappers.""" + tools: list[BaseTool] = [] + for wrapper in wrappers: + if isinstance(wrapper, CompositeToolWrapper): + comp_tools = wrapper.composite_tools_fn_getter(ctx) + for t in comp_tools: + tools.append(trace_langchain_tool(t, ctx)) + continue + + t = wrapper.tool_fn_getter(ctx) + tools.append(trace_langchain_tool(t, ctx)) + return tools + + +def get_tool_by_name(name: str, tools: list[BaseTool]) -> BaseTool | None: + """Get a tool by name, stripping any prefixes like 'default_api:'.""" + normalized_name = name + if ":" in normalized_name: + normalized_name = normalized_name.split(":")[-1] + + for t in tools: + t_name = t.name + if ":" in t_name: + t_name = t_name.split(":")[-1] + if t_name == normalized_name: + return t + return None diff --git a/artemis/tools/log_tool.py b/artemis/tools/log_tool.py new file mode 100644 index 0000000..13add08 --- /dev/null +++ b/artemis/tools/log_tool.py @@ -0,0 +1,108 @@ +# 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. + +"""Universal Log Analyzer tool for inspecting and analyzing Android system logs.""" + +from typing import Any + +from langchain_core.tools import BaseTool +from pydantic import BaseModel, Field + +from artemis.agents.log_analyzer.log_analyzer import LogAnalyzerNode +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, ToolRegistry +from artemis.tools.tool_wrapper import ToolWrapper +from artemis.tools.types import CyFunctionDetector +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +# pylint: disable=too-few-public-methods +class AnalyzeLogsArgs(BaseModel): + """Arguments schema for invoking the Log Analyzer subagent.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + specific_query: str = Field( + ..., + description=( + "Instructions for the log analyzer, specifying what to search for or analyze." + ), + ) + + +class AnalyzeLogsTool(ArtemisTool): + """Universal tool for analyzing Android logs using a dedicated hybrid log analyzer agent.""" + + def __init__(self): + super().__init__( + name="analyze_logs", + description=( + "[DIAGNOSTIC] Analyzes Android logs using a dedicated hybrid log analyzer agent. " + "Use this when you need to investigate failures, check for specific error " + "logs, or understand system events." + ), + args_schema=AnalyzeLogsArgs, + category="custom", + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + ctx: ArtemisContext | None = None, + specific_query: str | None = None, + state: State | None = None, + **kwargs: Any, + ) -> Any: + query = ( + specific_query + if specific_query is not None + else ( + kwargs.get("specific_query") + or kwargs.get("SpecificQuery") + or kwargs.get("query") + or kwargs.get("Query") + or "" + ) + ) + st = state if state is not None else kwargs.get("state") + logger.info(f"analyze_logs tool called with query: '{query}'") + if ctx is None: + return "Error: ArtemisContext is required for LogAnalyzer." + analyst = LogAnalyzerNode(ctx) + result = await analyst.run(query, st) + return result + + +# Universal tool instance & aliases +analyze_logs = AnalyzeLogsTool() +AnalyzeLogs = AnalyzeLogsTool +LogTool = AnalyzeLogsTool +ToolRegistry.register(analyze_logs) + + +def get_analyze_logs_tool(ctx: ArtemisContext) -> BaseTool: + """Exports analyze_logs as a LangChain BaseTool.""" + return trace_langchain_tool(analyze_logs.to_langchain_tool(ctx), ctx) + + +analyze_logs_wrapper = ToolWrapper( + tool_fn_getter=get_analyze_logs_tool, + on_success_fn=lambda output: f"Log Analyzer replied:\n{output}", + on_failure_fn=lambda error: f"Log Analyzer failed: {error}", +) diff --git a/artemis/tools/memory/notes.py b/artemis/tools/memory/notes.py new file mode 100644 index 0000000..5f81ef0 --- /dev/null +++ b/artemis/tools/memory/notes.py @@ -0,0 +1,66 @@ +# 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. + +"""Memory and Scratchpad Note Management tools.""" + +from typing import Any + +from pydantic import BaseModel, Field + +from artemis.tools.base import artemis_tool + +_IN_MEMORY_NOTES: dict[str, str] = {} + + +class ReadNoteArgs(BaseModel): + """Arguments schema for reading a note.""" + + title: str = Field(..., description="Note title to read.") + + +class SaveNoteArgs(BaseModel): + """Arguments schema for saving a note.""" + + title: str = Field(..., description="Note title to create/save.") + content: str = Field(..., description="Content text of the note.") + + +@artemis_tool( + name="read_note", + description="[MEMORY] Reads content of an existing scratchpad note.", + args_schema=ReadNoteArgs, + category="memory", +) +async def read_note( + title: str, + ctx: Any = None, # pylint: disable=unused-argument +) -> str: + """Reads content of an existing scratchpad note.""" + return _IN_MEMORY_NOTES.get(title, f"Note '{title}' not found.") + + +@artemis_tool( + name="save_note", + description="[MEMORY] Saves content into the scratchpad notes system.", + args_schema=SaveNoteArgs, + category="memory", +) +async def save_note( + title: str, + content: str, + ctx: Any = None, # pylint: disable=unused-argument +) -> str: + """Saves content into the scratchpad notes system.""" + _IN_MEMORY_NOTES[title] = content + return f"Note '{title}' saved successfully." diff --git a/artemis/tools/mobile/exec_tools.py b/artemis/tools/mobile/exec_tools.py new file mode 100644 index 0000000..6f0350e --- /dev/null +++ b/artemis/tools/mobile/exec_tools.py @@ -0,0 +1,730 @@ +# 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 asyncio +from typing import Any, Literal + +from langchain_core.messages import ToolMessage +from langchain_core.tools import BaseTool +from langgraph.types import Command +from pydantic import BaseModel, Field + +from artemis.constants import VALIDATOR_MESSAGES_KEY +from artemis.context import ArtemisContext +from artemis.controllers.unified_controller import UnifiedMobileController +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.tools.types import CyFunctionDetector +from artemis.utils.coordinates import parse_swipe_parameters +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +class ClickArgs(BaseModel): + """Arguments schema for clicking on normalized coordinates.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + target: list[int] = Field( + ..., + description=("Click target normalized coordinates [x, y] in 0-1000 scale."), + ) + + +class LongPressArgs(BaseModel): + """Arguments schema for long pressing on normalized coordinates.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + target: list[int] = Field( + ..., + description=("Long press target normalized coordinates [x, y] in 0-1000 scale."), + ) + duration: int = Field(1000, description="Long press duration in milliseconds (default 1000).") + + +class InputTextArgs(BaseModel): + """Arguments schema for inputting text into normalized coordinates.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + text: str = Field( + ..., + description="The text content to input. Supports multi-line strings with '\\n'.", + ) + target: list[int] = Field( + ..., + description=("Input target field normalized coordinates [x, y] in 0-1000 scale."), + ) + clear_exist: bool = Field( + True, + description=( + "Whether to clear existing text before typing (True to replace all" + " text, False to append at the end of existing content)." + ), + ) + + +class SwipeArgs(BaseModel): + """Arguments schema for performing swipe gestures.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + direction: Literal["up", "down", "left", "right"] | None = Field( + None, + description=( + "Smart directional swipe ('up', 'down', 'left', 'right') for general page" + " scrolling/browsing. Automatically applies safe vectors and a ~40% visual" + " overlap anchor: 'up' (drags bottom-to-top, scrolling down to reveal content below)," + " 'down' (drags top-to-bottom, scrolling up to reveal content above)," + " 'left' (drags right-to-left, scrolling right)," + " 'right' (drags left-to-right, scrolling left)." + ), + ) + start: list[int] | None = Field( + None, + description=( + "Start normalized coordinates [start_x, start_y] in 0-1000 scale for precise," + " local interactions (e.g. adjusting sliders, SeekBars, fine range selection, or drag-and-drop)." + ), + ) + end: list[int] | None = Field( + None, + description=( + "End normalized coordinates [end_x, end_y] in 0-1000 scale for precise," + " local interactions (e.g. adjusting sliders, SeekBars, fine range selection, or drag-and-drop)." + ), + ) + action: Literal["up", "down", "left", "right"] | list[int] | None = Field( + None, + description=( + "Backward-compatible swipe gesture: smart direction string ('up', 'down', 'left', 'right')" + " OR precise custom coordinates [start_x, start_y, end_x, end_y] in 0-1000 scale." + ), + ) + duration: int = Field( + 400, + description=( + "Optional swipe/drag duration in milliseconds (default 400). For drag-and-drop," + " list reordering, or sliding/adjusting sliders (e.g., volume, brightness, SeekBars)," + " set duration >= 1000 (e.g. 1500). If omitted for directional swipe, duration is computed automatically." + ), + ) + + +class PressKeyArgs(BaseModel): + """Arguments schema for pressing system keys.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + key: Literal["ENTER", "BACK", "HOME", "APP_SWITCH"] = Field( + ..., + description=("Standard Android system button name (ENTER, BACK, HOME, APP_SWITCH)."), + ) + + +CLICK_DOCSTRING = "[ACTION] Click on the target normalized coordinates on the screen." + + +class ClickTool(ArtemisTool): + """Universal tool for clicking on target normalized coordinates on the screen.""" + + def __init__(self, category: ToolCategory = "action"): + super().__init__( + name="click", + description=CLICK_DOCSTRING, + args_schema=ClickArgs, + category=category, + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals + async def execute( + self, + driver: BaseDeviceDriver | None = None, + ctx: ArtemisContext | None = None, + target: list[int] | None = None, + tool_call_id: str | None = None, + state: State | None = None, + **kwargs: Any, + ) -> Any: + tgt = target if target is not None else (kwargs.get("target") or kwargs.get("Target") or []) + 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") + + width = 1080 + height = 2400 + if ctx and hasattr(ctx, "device") and ctx.device: + width = getattr(ctx.device, "device_width", 1080) + height = getattr(ctx.device, "device_height", 2400) + elif driver and hasattr(driver, "screen_size"): + width, height = driver.screen_size + + try: + if not tgt or len(tgt) < 2: + raise ValueError(f"Invalid target coordinates: {tgt}") + nx, ny = tgt[0], tgt[1] + x = int(max(0, min(width - 1, nx * width / 1000))) + y = int(max(0, min(height - 1, ny * height / 1000))) + + if driver is not None and hasattr(driver, "tap"): + success = await driver.tap(x, y) + outcome = ( + f"Clicked at [{nx}, {ny}] (normalized) successfully." + if success + else f"Click failed at {tgt}" + ) + elif ctx is not None: + controller = UnifiedMobileController(ctx) + result = await controller.tap_at(x, y) + success = result.error is None + outcome = ( + f"Clicked at [{nx}, {ny}] (normalized) successfully." + if success + else f"Failed to click: {result.error}" + ) + else: + success = False + outcome = "Error during click: No driver or context provided." + except Exception as e: # pylint: disable=broad-exception-caught + success = False + outcome = f"Error during click: {e}" + + if st and callable(getattr(st, "asanitize_update", None)): + tool_message = ToolMessage( + tool_call_id=tcid or "", + content=outcome, + status="success" if success else "error", + ) + return Command( + update=await st.asanitize_update( + ctx=ctx, + update={VALIDATOR_MESSAGES_KEY: [tool_message]}, + agent="validator", + ) + ) + + return outcome + + +# Universal tool instance & aliases +click = ClickTool() +Click = ClickTool +ToolRegistry.register(click) + + +def get_click_tool(ctx: ArtemisContext) -> BaseTool: + """Exports click as a LangChain BaseTool.""" + return trace_langchain_tool(click.to_langchain_tool(ctx), ctx) + + +LONG_PRESS_DOCSTRING = "[ACTION] Long press on the target normalized coordinates on the screen." + + +class LongPressTool(ArtemisTool): + """Universal tool for long pressing on target normalized coordinates on the screen.""" + + def __init__(self, category: ToolCategory = "action"): + super().__init__( + name="long_press", + description=LONG_PRESS_DOCSTRING, + args_schema=LongPressArgs, + category=category, + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals + async def execute( + self, + driver: BaseDeviceDriver | None = None, + ctx: ArtemisContext | None = None, + target: list[int] | None = None, + duration: int = 1000, + tool_call_id: str | None = None, + state: State | None = None, + **kwargs: Any, + ) -> Any: + tgt = target if target is not None else (kwargs.get("target") or kwargs.get("Target") or []) + dur = ( + duration + if duration is not None + else (kwargs.get("duration") or kwargs.get("Duration") or 1000) + ) + 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") + + width = 1080 + height = 2400 + if ctx and hasattr(ctx, "device") and ctx.device: + width = getattr(ctx.device, "device_width", 1080) + height = getattr(ctx.device, "device_height", 2400) + elif driver and hasattr(driver, "screen_size"): + width, height = driver.screen_size + + try: + if not tgt or len(tgt) < 2: + raise ValueError(f"Invalid target coordinates: {tgt}") + nx, ny = tgt[0], tgt[1] + x = int(max(0, min(width - 1, nx * width / 1000))) + y = int(max(0, min(height - 1, ny * height / 1000))) + + if driver is not None and hasattr(driver, "long_press"): + success = await driver.long_press(x, y, duration_ms=dur) + outcome = ( + f"Long pressed at [{nx}, {ny}] (normalized) successfully." + if success + else f"Failed to long press: Long press failed at {tgt}" + ) + elif ctx is not None: + controller = UnifiedMobileController(ctx) + result = await controller.tap_at(x, y, long_press=True, long_press_duration=dur) + success = result.error is None + outcome = ( + f"Long pressed at [{nx}, {ny}] (normalized) successfully." + if success + else f"Failed to long press: {result.error}" + ) + else: + success = False + outcome = "Error during long press: No driver or context provided." + except Exception as e: # pylint: disable=broad-exception-caught + success = False + outcome = f"Error during long press: {e}" + + if st and callable(getattr(st, "asanitize_update", None)): + tool_message = ToolMessage( + tool_call_id=tcid or "", + content=outcome, + status="success" if success else "error", + ) + return Command( + update=await st.asanitize_update( + ctx=ctx, + update={VALIDATOR_MESSAGES_KEY: [tool_message]}, + agent="validator", + ) + ) + + return outcome + + +# Universal tool instance & aliases +long_press = LongPressTool() +LongPress = LongPressTool +ToolRegistry.register(long_press) + + +def get_long_press_tool(ctx: ArtemisContext) -> BaseTool: + """Exports long_press as a LangChain BaseTool.""" + return trace_langchain_tool(long_press.to_langchain_tool(ctx), ctx) + + +INPUT_TEXT_DOCSTRING = "[ACTION] Type text into the target input field coordinates." + + +class InputTextTool(ArtemisTool): + """Universal tool for typing text into target input field coordinates.""" + + def __init__(self, category: ToolCategory = "action"): + super().__init__( + name="input_text", + description=INPUT_TEXT_DOCSTRING, + args_schema=InputTextArgs, + category=category, + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals + async def execute( + self, + driver: BaseDeviceDriver | None = None, + ctx: ArtemisContext | None = None, + text: str | None = None, + target: list[int] | None = None, + clear_exist: bool = True, + tool_call_id: str | None = None, + state: State | None = None, + **kwargs: Any, + ) -> Any: + txt = text if text is not None else (kwargs.get("text") or kwargs.get("Text") or "") + tgt = target if target is not None else (kwargs.get("target") or kwargs.get("Target") or []) + clr = ( + clear_exist + if clear_exist is not None + else kwargs.get("clear_exist", kwargs.get("ClearExist", True)) + ) + 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") + + width = 1080 + height = 2400 + if ctx and hasattr(ctx, "device") and ctx.device: + width = getattr(ctx.device, "device_width", 1080) + height = getattr(ctx.device, "device_height", 2400) + elif driver and hasattr(driver, "screen_size"): + width, height = driver.screen_size + + try: + if not tgt or len(tgt) < 2: + raise ValueError(f"Invalid target coordinates: {tgt}") + nx, ny = tgt[0], tgt[1] + x = int(max(0, min(width - 1, nx * width / 1000))) + y = int(max(0, min(height - 1, ny * height / 1000))) + + if driver is not None and hasattr(driver, "input_text"): + await driver.tap(x, y) + await asyncio.sleep(0.2) + success = await driver.input_text(txt, clear_existing=clr) + outcome = ( + f"Typed '{txt}' successfully." + if success + else f"Failed to input text into {tgt}" + ) + elif ctx is not None: + controller = UnifiedMobileController(ctx) + tap_result = await controller.tap_at(x, y) + if tap_result.error: + raise RuntimeError(f"Failed to focus input field: {tap_result.error}") + + if clr: + await controller.erase_text() + else: + await controller.press_key("123") # Move to end + + success = await controller.type_text(txt, clear_existing=False) + outcome = f"Typed '{txt}' successfully." if success else "Failed to type text." + else: + success = False + outcome = "Error during input text: No driver or context provided." + except Exception as e: # pylint: disable=broad-exception-caught + success = False + outcome = f"Error during input text: {e}" + + if st and callable(getattr(st, "asanitize_update", None)): + tool_message = ToolMessage( + tool_call_id=tcid or "", + content=outcome, + status="success" if success else "error", + ) + return Command( + update=await st.asanitize_update( + ctx=ctx, + update={VALIDATOR_MESSAGES_KEY: [tool_message]}, + agent="validator", + ) + ) + + return outcome + + +# Universal tool instance & aliases +input_text = InputTextTool() +InputText = InputTextTool +ToolRegistry.register(input_text) + + +def get_input_text_tool(ctx: ArtemisContext) -> BaseTool: + """Exports input_text as a LangChain BaseTool.""" + return trace_langchain_tool(input_text.to_langchain_tool(ctx), ctx) + + +SWIPE_DOCSTRING = ( + "[ACTION] Perform a swipe, drag, or slider-adjustment gesture on the" + " screen.\n\n" + "โ€ข Smart Directional Scrolling ('direction'): Best for general browsing" + " and page scrolling. Automatically computes safe swipe vectors (locking" + " axes to prevent edge/sidebar mis-clicks), retains a ~40% visual overlap" + " anchor for zero-omission traversal, and prevents inertial flings.\n" + "โ€ข Precise Coordinate Gestures ('start', 'end'): Best for local," + " fine-grained interactions such as adjusting sliders/SeekBars (e.g.," + " volume, brightness, progress bars) or drag-and-drop / list reordering." + " Always drag slightly PAST the target position to overcome touch slop" + " and reliably trigger the update. When setting a slider to Maximum (100%)" + " or Minimum (0%), swipe fully to the extreme boundary." +) + + +class SwipeTool(ArtemisTool): + """Universal tool for performing swipe gestures on the screen.""" + + def __init__(self, category: ToolCategory = "action"): + super().__init__( + name="swipe", + description=SWIPE_DOCSTRING, + args_schema=SwipeArgs, + category=category, + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-branches,too-many-locals,too-many-statements + async def execute( + self, + driver: BaseDeviceDriver | None = None, + ctx: ArtemisContext | None = None, + action: Literal["up", "down", "left", "right"] | list[int] | str | None = None, + direction: Literal["up", "down", "left", "right"] | str | None = None, + start: list[int] | None = None, + end: list[int] | None = None, + duration: int = 400, + tool_call_id: str | None = None, + state: State | None = None, + **kwargs: Any, + ) -> Any: + params = dict(kwargs) + if action is not None: + params["action"] = action + if direction is not None: + params["direction"] = direction + if start is not None: + params["start"] = start + if end is not None: + params["end"] = end + if duration is not None: + params["duration"] = duration + + kind, target, dur = parse_swipe_parameters(params, default_duration=duration or 400) + + 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") + + width = 1080 + height = 2400 + if ctx and hasattr(ctx, "device") and ctx.device: + width = getattr(ctx.device, "device_width", 1080) + height = getattr(ctx.device, "device_height", 2400) + elif driver and hasattr(driver, "screen_size"): + width, height = driver.screen_size + + try: + if kind is None or target is None: + raise ValueError( + f"Invalid swipe action parameter: {action or direction or (start, end) or kwargs}" + ) + + if ( + driver is not None + and hasattr(driver, "swipe_direction") + and hasattr(driver, "swipe") + ): + if kind == "direction" and isinstance(target, str): + success = await driver.swipe_direction(target, duration_ms=dur) + outcome = ( + "Swipe completed successfully." if success else f"Swipe {target} failed" + ) + elif kind == "coords" and isinstance(target, list) and len(target) == 4: + start_x = int(max(0, min(width - 1, target[0] * width / 1000))) + start_y = int(max(0, min(height - 1, target[1] * height / 1000))) + end_x = int(max(0, min(width - 1, target[2] * width / 1000))) + end_y = int(max(0, min(height - 1, target[3] * height / 1000))) + success = await driver.swipe(start_x, start_y, end_x, end_y, duration_ms=dur) + outcome = ( + "Swipe completed successfully." + if success + else f"Swipe coordinates {target} failed" + ) + else: + raise ValueError(f"Invalid swipe action parameter: {target}") + elif ctx is not None: + controller = UnifiedMobileController(ctx) + + if kind == "direction" and isinstance(target, str): + if target == "up": + coords = [ + width // 2, + int(height * 0.8), + width // 2, + int(height * 0.2), + ] + elif target == "down": + coords = [ + width // 2, + int(height * 0.2), + width // 2, + int(height * 0.8), + ] + elif target == "left": + coords = [ + int(width * 0.8), + height // 2, + int(width * 0.2), + height // 2, + ] + elif target == "right": + coords = [ + int(width * 0.2), + height // 2, + int(width * 0.8), + height // 2, + ] + else: + raise ValueError(f"Invalid direction: {target}") + + start_x, start_y, end_x, end_y = coords + elif kind == "coords" and isinstance(target, list) and len(target) == 4: + nx1, ny1, nx2, ny2 = target + start_x = int(max(0, min(width - 1, nx1 * width / 1000))) + start_y = int(max(0, min(height - 1, ny1 * height / 1000))) + end_x = int(max(0, min(width - 1, nx2 * width / 1000))) + end_y = int(max(0, min(height - 1, ny2 * height / 1000))) + else: + raise ValueError(f"Invalid swipe action parameter: {target}") + + error = await controller.swipe_coords(start_x, start_y, end_x, end_y, dur) + success = error is None + outcome = "Swipe completed successfully." if success else f"Swipe failed: {error}" + else: + success = False + outcome = "Error during swipe: No driver or context provided." + except Exception as e: # pylint: disable=broad-exception-caught + success = False + outcome = f"Error during swipe: {e}" + + if st and callable(getattr(st, "asanitize_update", None)): + tool_message = ToolMessage( + tool_call_id=tcid or "", + content=outcome, + status="success" if success else "error", + ) + return Command( + update=await st.asanitize_update( + ctx=ctx, + update={VALIDATOR_MESSAGES_KEY: [tool_message]}, + agent="validator", + ) + ) + + return outcome + + +# Universal tool instance & aliases +swipe = SwipeTool() +Swipe = SwipeTool +ToolRegistry.register(swipe) + + +def get_swipe_tool(ctx: ArtemisContext) -> BaseTool: + """Exports swipe as a LangChain BaseTool.""" + return trace_langchain_tool(swipe.to_langchain_tool(ctx), ctx) + + +PRESS_KEY_DOCSTRING = "[ACTION] Press a physical or virtual system button." + + +class PressKeyTool(ArtemisTool): + """Universal tool for pressing a physical or virtual system button.""" + + def __init__(self, category: ToolCategory = "action"): + super().__init__( + name="press_key", + description=PRESS_KEY_DOCSTRING, + args_schema=PressKeyArgs, + category=category, + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments + async def execute( + self, + driver: BaseDeviceDriver | None = None, + ctx: ArtemisContext | None = None, + key: Literal["ENTER", "BACK", "HOME", "APP_SWITCH"] | str | None = None, + tool_call_id: str | None = None, + state: State | None = None, + **kwargs: Any, + ) -> Any: + k = key if key is not None else (kwargs.get("key") or kwargs.get("Key") or "") + 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") + + try: + if not k: + raise ValueError("Key parameter is required.") + + if driver is not None and hasattr(driver, "press_key"): + success = await driver.press_key(k) + outcome = ( + f"Pressed key '{k}' successfully." if success else f"Failed to press key '{k}'." + ) + elif ctx is not None: + controller = UnifiedMobileController(ctx) + success = False + if k == "ENTER": + success = await controller.press_enter() + elif k == "BACK": + success = await controller.go_back() + elif k == "HOME": + success = await controller.go_home() + elif k == "APP_SWITCH": + success = await controller.press_key("KEYCODE_APP_SWITCH") + else: + success = await controller.press_key(k) + + outcome = ( + f"Pressed key '{k}' successfully." if success else f"Failed to press key '{k}'." + ) + else: + success = False + outcome = "Error during press key: No driver or context provided." + except Exception as e: # pylint: disable=broad-exception-caught + success = False + outcome = f"Error during press key: {e}" + + if st and callable(getattr(st, "asanitize_update", None)): + tool_message = ToolMessage( + tool_call_id=tcid or "", + content=outcome, + status="success" if success else "error", + ) + return Command( + update=await st.asanitize_update( + ctx=ctx, + update={VALIDATOR_MESSAGES_KEY: [tool_message]}, + agent="validator", + ) + ) + + return outcome + + +# Universal tool instance & aliases +press_key = PressKeyTool() +PressKey = PressKeyTool +ToolRegistry.register(press_key) + + +def get_press_key_tool(ctx: ArtemisContext) -> BaseTool: + """Exports press_key as a LangChain BaseTool.""" + return trace_langchain_tool(press_key.to_langchain_tool(ctx), ctx) + + +# Wrappers to expose +click_wrapper = ToolWrapper( + tool_fn_getter=get_click_tool, + on_success_fn=lambda *args, **kwargs: "Success", + on_failure_fn=lambda *args, **kwargs: "Failure", +) +long_press_wrapper = ToolWrapper( + tool_fn_getter=get_long_press_tool, + on_success_fn=lambda *args, **kwargs: "Success", + on_failure_fn=lambda *args, **kwargs: "Failure", +) +input_text_wrapper = ToolWrapper( + tool_fn_getter=get_input_text_tool, + on_success_fn=lambda *args, **kwargs: "Success", + on_failure_fn=lambda *args, **kwargs: "Failure", +) +swipe_wrapper = ToolWrapper( + tool_fn_getter=get_swipe_tool, + on_success_fn=lambda *args, **kwargs: "Success", + on_failure_fn=lambda *args, **kwargs: "Failure", +) +press_key_wrapper = ToolWrapper( + tool_fn_getter=get_press_key_tool, + on_success_fn=lambda *args, **kwargs: "Success", + on_failure_fn=lambda *args, **kwargs: "Failure", +) diff --git a/artemis/tools/mobile/launch_app.py b/artemis/tools/mobile/launch_app.py new file mode 100644 index 0000000..1f0db1c --- /dev/null +++ b/artemis/tools/mobile/launch_app.py @@ -0,0 +1,215 @@ +# 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. + +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, Field + +from artemis.agents.hopper.hopper import HopperOutput, hopper +from artemis.constants import VALIDATOR_MESSAGES_KEY +from artemis.context import ArtemisContext +from artemis.controllers.platform_specific_commands_controller import ( + list_packages_async, +) +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.tools.types import CyFunctionDetector +from artemis.utils.app_launch_utils import launch_app_with_retries +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +class LaunchAppArgs(BaseModel): + """Arguments schema for launching an application.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + app_name: str = Field( + ..., + description="The natural language name of the application to launch.", + ) + + +LAUNCH_APP_DOCSTRING = ( + "[ACTION] Finds and launches an application on the device using its natural language name." +) + + +async def find_package(ctx: ArtemisContext, app_name: str, use_fallback: bool = True) -> str | None: + """Finds the package name for a given application name. + + Returns None if package not found or on error. + """ + package_cache = getattr(ctx, "package_cache", None) + if package_cache is None or not isinstance(package_cache, dict): + try: + ctx.package_cache = {} + package_cache = ctx.package_cache + except Exception: # pylint: disable=broad-exception-caught + package_cache = {} + + if isinstance(package_cache, dict) and app_name in package_cache: + logger.info(f"Cache hit for app '{app_name}': {package_cache[app_name]}") + return package_cache[app_name] + + try: + all_packages = await list_packages_async(ctx=ctx) + package_set = {p.strip() for p in all_packages.split("\n") if p.strip()} + + # Fast path: If app_name is already directly an installed package name + if app_name in package_set: + if isinstance(package_cache, dict): + package_cache[app_name] = app_name + return app_name + + hopper_output: HopperOutput = await hopper( + ctx=ctx, + request=(f"I'm looking for the package name of the following app: '{app_name}'"), + data=all_packages, + use_fallback=use_fallback, + ) + if not hopper_output.found or not hopper_output.output: + if isinstance(package_cache, dict): + package_cache[app_name] = None + return None + + package_name = hopper_output.output.strip() + if package_name not in package_set: + logger.warning( + f"Hopper returned package '{package_name}' for '{app_name}', " + "but it is NOT physically installed on the device!" + ) + if isinstance(package_cache, dict): + package_cache[app_name] = None + return None + + if isinstance(package_cache, dict): + package_cache[app_name] = package_name + return package_name + except Exception as e: # pylint: disable=broad-exception-caught + logger.error(f"Failed to find package for '{app_name}': {e}") + return None + + +class LaunchAppTool(ArtemisTool): + """Universal tool for finding and launching an application on the device.""" + + def __init__(self, category: ToolCategory = "action"): + super().__init__( + name="launch_app", + description=LAUNCH_APP_DOCSTRING, + args_schema=LaunchAppArgs, + category=category, + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals + async def execute( + self, + driver: BaseDeviceDriver | None = None, + ctx: ArtemisContext | None = None, + app_name: str | None = None, + tool_call_id: str | None = None, + state: State | None = None, + **kwargs: Any, + ) -> Any: + app = ( + app_name + if app_name is not None + else (kwargs.get("app_name") or kwargs.get("AppName") or "") + ) + 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") + + success = False + error_msg = None + + try: + if not app: + raise ValueError("app_name parameter is required.") + + if ctx is not None: + package_name = await find_package(ctx=ctx, app_name=app) + if not package_name: + success = False + outcome = f"Failed to launch app '{app}': Package not found." + error_msg = "Package not found." + else: + success, error_msg = await launch_app_with_retries( + ctx=ctx, app_package=package_name + ) + outcome = ( + f"App '{app}' launched successfully." + if success + else f"Failed to launch app '{app}': {error_msg}" + ) + elif driver is not None and hasattr(driver, "launch_app"): + success = await driver.launch_app(app) + outcome = ( + f"App '{app}' launched successfully." + if success + else f"Failed to launch app '{app}'." + ) + error_msg = None if success else "Launch failed." + else: + success = False + outcome = "Error during launch app: No driver or context provided." + error_msg = "No driver or context provided." + except Exception as e: # pylint: disable=broad-exception-caught + success = False + outcome = f"Error during launch app: {e}" + error_msg = str(e) + + if st and callable(getattr(st, "asanitize_update", None)): + additional_kwargs = {} if success else ({"error": error_msg} if error_msg else {}) + tool_message = ToolMessage( + tool_call_id=tcid or "", + content=outcome, + additional_kwargs=additional_kwargs, + status="success" if success else "error", + ) + return Command( + update=await st.asanitize_update( + ctx=ctx, + update={ + VALIDATOR_MESSAGES_KEY: [tool_message], + }, + agent="validator", + ) + ) + + return outcome + + +# Universal tool instance & aliases +launch_app = LaunchAppTool() +LaunchApp = LaunchAppTool +ToolRegistry.register(launch_app) + + +def get_launch_app_tool(ctx: ArtemisContext) -> BaseTool: + """Exports launch_app as a LangChain BaseTool.""" + return trace_langchain_tool(launch_app.to_langchain_tool(ctx), ctx) + + +launch_app_wrapper = ToolWrapper( + tool_fn_getter=get_launch_app_tool, + on_success_fn=lambda *args, **kwargs: "Success", + on_failure_fn=lambda *args, **kwargs: "Failure", +) diff --git a/artemis/tools/mobile/log_utils.py b/artemis/tools/mobile/log_utils.py new file mode 100644 index 0000000..39d5e4f --- /dev/null +++ b/artemis/tools/mobile/log_utils.py @@ -0,0 +1,138 @@ +# 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. + +from datetime import datetime, timedelta +import time +from artemis.context import ArtemisContext +from artemis.controllers.platform_specific_commands_controller import get_adb_device +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +def parse_log_time(time_str: str) -> datetime: + """Parses a log timestamp string into a datetime object.""" + now = datetime.now() + current_year = now.year + + def _try_parse(year: int) -> datetime | None: + for fmt in ("%Y-%m-%d %H:%M:%S.%f", "%Y-%m-%d %H:%M:%S"): + try: + return datetime.strptime(f"{year}-{time_str}", fmt) + except ValueError: + pass + return None + + dt = _try_parse(current_year) + if dt is None: + raise ValueError(f"Unsupported time format: {time_str}") + + # If parsed time is in the future (allowing 1 hour grace for clock drift), + # we assume it belongs to the previous year. + if dt > now + timedelta(hours=1): + prev_dt = _try_parse(current_year - 1) + if prev_dt is not None: + dt = prev_dt + + return dt + + +def resolve_time( + time_input: str | float | None, session_start: float | None +) -> tuple[str | None, datetime | None]: + """Resolves relative or absolute time input into formatted string and datetime.""" + if time_input is None: + return None, None + + if isinstance(time_input, (int, float)) or ( + isinstance(time_input, str) and time_input.replace(".", "", 1).isdigit() + ): + # It's relative seconds + if session_start is None: + raise ValueError("Cannot resolve relative time without session start time.") + rel_sec = float(time_input) + abs_ts = session_start + rel_sec + + t_struct = time.localtime(abs_ts) + t_str = time.strftime("%m-%d %H:%M:%S", t_struct) + ms = int((abs_ts - int(abs_ts)) * 1000) + formatted_str = f"{t_str}.{ms:03d}" + + dt = datetime.fromtimestamp(abs_ts) + return formatted_str, dt + + # Assume it's absolute time string "MM-DD HH:MM:SS.ms" + try: + dt = parse_log_time(str(time_input)) + except ValueError as e: + raise ValueError(f"Unsupported time format: {time_input}") from e + return str(time_input), dt + + +# pylint: disable=too-many-locals +def fetch_and_filter_logs( + ctx: ArtemisContext, + lines: int = 1000, + since_time: str | None = None, + until_time: str | None = None, +) -> str: + """Fetches logs from ADB device and filters them by time if needed.""" + device = get_adb_device(ctx) + + session_start = ctx.data_engine.session_start_time if ctx.data_engine else None + + since_str, since_dt = resolve_time(since_time, session_start) + _, until_dt = resolve_time(until_time, session_start) + + if since_str: + try: + logger.info(f"Attempting to fetch logs since {since_str}...") + logs = device.shell(f'logcat -v threadtime -t "{since_str}"') + except Exception as e: # pylint: disable=broad-exception-caught + logger.warning( + f"Failed to fetch logs since time directly: {e}. Falling back" + " to reading lines and filtering." + ) + logs = device.shell(f"logcat -v threadtime -t {lines}") + else: + logger.info(f"Reading last {lines} lines of logs...") + logs = device.shell(f"logcat -v threadtime -t {lines}") + + if since_dt or until_dt: + filtered_lines = [] + for line in logs.splitlines(): + if len(line) >= 18: + log_time_str = line[:18] + try: + log_dt = parse_log_time(log_time_str) + + if since_dt and log_dt < since_dt: + continue + if until_dt and log_dt > until_dt: + continue + + filtered_lines.append(line) + except ValueError: + pass + # Cap to the requested line limit so timestamp fetches don't overflow context + if 0 < lines < len(filtered_lines): + filtered_lines = filtered_lines[-lines:] + logs = "\n".join(filtered_lines) + else: + # Ensure fallback capping even when no datetime parsing occurred + log_lines = logs.splitlines() + if 0 < lines < len(log_lines): + logs = "\n".join(log_lines[-lines:]) + + return logs diff --git a/artemis/tools/mobile/ocr.py b/artemis/tools/mobile/ocr.py new file mode 100644 index 0000000..f06fc97 --- /dev/null +++ b/artemis/tools/mobile/ocr.py @@ -0,0 +1,210 @@ +# 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}", +) diff --git a/artemis/tools/mobile/read_hierarchy.py b/artemis/tools/mobile/read_hierarchy.py new file mode 100644 index 0000000..f05d4ed --- /dev/null +++ b/artemis/tools/mobile/read_hierarchy.py @@ -0,0 +1,106 @@ +# 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 inspect +from typing import Any + +from langchain_core.tools import BaseTool +from pydantic import BaseModel + +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.tools.types import CyFunctionDetector +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +class GetUiHierarchyArgs(BaseModel): + """Arguments schema for getting UI hierarchy.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + + +GET_UI_HIERARCHY_DOCSTRING = ( + "[DIAGNOSTIC] Retrieves the current screen UI hierarchy XML using a persistent" + " UI Automator client.\n\nUse this instead of running 'uiautomator dump'" + " command via adb shell, as it is much faster and more stable." +) + + +class GetUiHierarchyTool(ArtemisTool): + """Universal tool for retrieving screen UI hierarchy XML.""" + + def __init__(self, category: ToolCategory = "diagnostic"): + super().__init__( + name="get_ui_hierarchy", + description=GET_UI_HIERARCHY_DOCSTRING, + args_schema=GetUiHierarchyArgs, + 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, # pylint: disable=unused-argument + state: State | None = None, # pylint: disable=unused-argument + **kwargs: Any, + ) -> str: + try: + logger.info("get_ui_hierarchy tool called.") + if ctx is not None and getattr(ctx, "ui_adb_client", None) is not None: + xml_hierarchy = ctx.ui_adb_client.get_hierarchy() + if inspect.iscoroutine(xml_hierarchy): + xml_hierarchy = await xml_hierarchy + return str(xml_hierarchy) + if driver is not None and hasattr(driver, "get_screen_data"): + screen_data = await driver.get_screen_data() + xml_hierarchy = getattr(screen_data, "ui_hierarchy_xml", None) + if xml_hierarchy: + return str(xml_hierarchy) + return "Error retrieving UI hierarchy: No UI hierarchy in screen data." + if driver is not None and hasattr(driver, "get_ui_hierarchy"): + xml_hierarchy = await driver.get_ui_hierarchy() + if inspect.iscoroutine(xml_hierarchy): + xml_hierarchy = await xml_hierarchy + return str(xml_hierarchy) + return "Error retrieving UI hierarchy: No UI Automator client or driver provided." + except Exception as e: # pylint: disable=broad-exception-caught + logger.error(f"Failed to get UI hierarchy: {e}") + return f"Error retrieving UI hierarchy: {e}" + + +# Universal tool instance & aliases +get_ui_hierarchy = GetUiHierarchyTool() +GetUiHierarchy = GetUiHierarchyTool +GetUIHierarchy = GetUiHierarchyTool +ToolRegistry.register(get_ui_hierarchy) + + +def get_ui_hierarchy_tool(ctx: ArtemisContext) -> BaseTool: + """Exports get_ui_hierarchy as a LangChain BaseTool.""" + return trace_langchain_tool(get_ui_hierarchy.to_langchain_tool(ctx), ctx) + + +ui_hierarchy_wrapper = ToolWrapper( + tool_fn_getter=get_ui_hierarchy_tool, + on_success_fn=lambda output: output, + on_failure_fn=lambda error: f"Failed to retrieve UI hierarchy: {error}", +) diff --git a/artemis/tools/mobile/read_logs.py b/artemis/tools/mobile/read_logs.py new file mode 100644 index 0000000..bfab9de --- /dev/null +++ b/artemis/tools/mobile/read_logs.py @@ -0,0 +1,116 @@ +# 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. + +from typing import Any + +from langchain_core.tools.base import BaseTool +from pydantic import BaseModel, Field + +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.mobile.log_utils import fetch_and_filter_logs +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +class ReadLogsArgs(BaseModel): + """Arguments schema for read_logs tool.""" + + lines: int | None = Field( + default=200, + description=( + "Number of recent log lines to read. Default is 200. Use larger" + " numbers only if necessary." + ), + ) + since_time: str | None = Field( + default=None, + description=( + "Target timestamp in 'MM-DD HH:MM:SS.ms' format OR relative seconds" + " (e.g., '15.5') to read logs since. If provided, we will try to" + " fetch logs starting from this time." + ), + ) + until_time: str | None = Field( + default=None, + description=( + "Target timestamp in 'MM-DD HH:MM:SS.ms' format OR relative seconds" + " (e.g., '20.0') to read logs until. If provided, we will filter" + " logs up to this time." + ), + ) + + +READ_LOGS_DOCSTRING = ( + "[DIAGNOSTIC] Fetches raw Android logs.\n\n" + "- Use when: You need to scan logs for a time window or see the sequence" + " of events.\n" + "- Input: lines (int, optional): Number of lines to read (default 200)." + " since_time (str, optional): Timestamp or relative seconds to read since." + " until_time (str, optional): Timestamp or relative seconds to read until.\n" + "- Output: Raw log text." +) + + +class ReadLogsTool(ArtemisTool): + """Universal tool for fetching raw Android logs.""" + + def __init__(self, category: ToolCategory = "diagnostic"): + super().__init__( + name="read_logs", + description=READ_LOGS_DOCSTRING, + args_schema=ReadLogsArgs, + category=category, + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + ctx: ArtemisContext | None = None, + lines: int | None = 200, + since_time: str | None = None, + until_time: str | None = None, + tool_call_id: str | None = None, # pylint: disable=unused-argument + state: State | None = None, # pylint: disable=unused-argument + **kwargs: Any, + ) -> str: + try: + num_lines = 200 if lines is None else lines + logs = fetch_and_filter_logs( + ctx=ctx, + lines=num_lines, + since_time=since_time, + until_time=until_time, + ) + return logs + except Exception as e: # pylint: disable=broad-exception-caught + logger.error(f"Failed to read logs: {e}") + return f"Failed to read logs: {e}" + + +# Universal tool instance & aliases +read_logs = ReadLogsTool() +ReadLogs = ReadLogsTool +ReadLogsToolAlias = ReadLogsTool +ToolRegistry.register(read_logs) + + +def get_read_logs_tool(ctx: ArtemisContext) -> BaseTool: + """Exports read_logs as a LangChain BaseTool.""" + return trace_langchain_tool(read_logs.to_langchain_tool(ctx), ctx) diff --git a/artemis/tools/mobile/search_logs.py b/artemis/tools/mobile/search_logs.py new file mode 100644 index 0000000..af328ef --- /dev/null +++ b/artemis/tools/mobile/search_logs.py @@ -0,0 +1,237 @@ +# 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 re +from typing import Any + +from langchain_core.tools.base import BaseTool +from pydantic import BaseModel, Field + +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.mobile.log_utils import fetch_and_filter_logs +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +# pylint: disable=too-many-branches,too-many-locals,too-many-statements +def search_and_merge_logs( + logs: str, + keyword: str, + context_lines: int = 50, + is_regex: bool = False, +) -> str: + """Searches for keyword/pattern in logs and merges overlapping context windows.""" + log_lines = logs.splitlines() + + if is_regex: + try: + pattern = re.compile(keyword, re.IGNORECASE) + except re.error as e: + return f"Error: Invalid regex: {e}" + else: + keyword_lower = keyword.lower() + + # 1. Collect all raw match ranges + ranges = [] + for i, line in enumerate(log_lines): + matched = False + if is_regex: + if pattern.search(line): + matched = True + else: + if keyword_lower in line.lower(): + matched = True + + if matched: + start = max(0, i - context_lines) + end = min(len(log_lines), i + context_lines + 1) + ranges.append((start, end, i)) + + if not ranges: + return f"No matches found for '{keyword}'." + + # 2. Merge overlapping/adjacent ranges + ranges.sort(key=lambda x: x[0]) + + merged_ranges = [] + curr_start, curr_end, curr_matches = ( + ranges[0][0], + ranges[0][1], + [ranges[0][2]], + ) + + for start, end, match_idx in ranges[1:]: + if start <= curr_end: # Overlaps or adjacent + curr_end = max(curr_end, end) + curr_matches.append(match_idx) + else: + merged_ranges.append((curr_start, curr_end, curr_matches)) + curr_start = start + curr_end = end + curr_matches = [match_idx] + merged_ranges.append((curr_start, curr_end, curr_matches)) + + # 3. Format output + matches = [] + total_lines = 0 + max_return_lines = 300 + + truncated = False + for start, end, match_indices in merged_ranges: + if total_lines >= max_return_lines: + truncated = True + break + + match_str = ", ".join(map(str, match_indices)) + matches.append(f"--- Match at line(s) {match_str} ---") + for idx in range(start, end): + line_content = log_lines[idx] + if idx in match_indices: + matches.append(f"{idx:4d}: [MATCH] {line_content}") + else: + matches.append(f"{idx:4d}: {line_content}") + + total_lines += 1 + if total_lines >= max_return_lines: + truncated = True + break + + matches.append("-" * 20) + + if truncated: + matches.append( + "\n... [SYSTEM WARNING: Output truncated to protect Context Window." + " Too many matches exceeded the 300 lines limit. You MUST use a" + " more specific keyword, narrow the time range, or use" + " context_lines=0.] ..." + ) + + return "\n".join(matches) + + +class SearchLogsArgs(BaseModel): + """Arguments schema for search_logs tool.""" + + keyword: str = Field( + ..., + description=("The keyword or regular expression to search for. Case-insensitive."), + ) + is_regex: bool = Field( + default=False, + description=("Whether to treat keyword as a regular expression. Default is False."), + ) + lines: int | None = Field( + default=10000, + description=("Number of recent log lines to search in. Default is 10000."), + ) + since_time: str | None = Field( + default=None, + description=( + "Target timestamp in 'MM-DD HH:MM:SS.ms' format OR relative seconds" + " (e.g., '15.5') to search logs since. If provided, we will search" + " logs starting from this time." + ), + ) + until_time: str | None = Field( + default=None, + description=( + "Target timestamp in 'MM-DD HH:MM:SS.ms' format OR relative seconds" + " (e.g., '20.0') to search logs until. If provided, we will filter" + " logs up to this time." + ), + ) + context_lines: int | None = Field( + default=0, + description=( + "Number of lines to include before and after each match for" + " context. Default is 0. Only increase if you need more context" + " around a very specific match." + ), + ) + + +SEARCH_LOGS_DOCSTRING = ( + "[DIAGNOSTIC] Searches logs for a keyword or regex pattern.\n\n" + "- Use when: Looking for specific errors, tags, or keywords without" + " reading all logs.\n" + "- Input: keyword (str): Search pattern. is_regex (bool, optional):" + " Treat keyword as regex. context_lines (int, optional): Lines around" + " match. since_time (str, optional): Timestamp or relative seconds to" + " search since. until_time (str, optional): Timestamp or relative seconds" + " to search until.\n" + "- Output: Matching lines with context." +) + + +class SearchLogsTool(ArtemisTool): + """Universal tool for searching logs for keywords or regex patterns.""" + + def __init__(self, category: ToolCategory = "diagnostic"): + super().__init__( + name="search_logs", + description=SEARCH_LOGS_DOCSTRING, + args_schema=SearchLogsArgs, + category=category, + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + ctx: ArtemisContext | None = None, + keyword: str = "", + is_regex: bool = False, + lines: int | None = 10000, + since_time: str | None = None, + until_time: str | None = None, + context_lines: int | None = 0, + tool_call_id: str | None = None, # pylint: disable=unused-argument + state: State | None = None, # pylint: disable=unused-argument + **kwargs: Any, + ) -> str: + try: + num_lines = 10000 if lines is None else lines + ctx_lines = 0 if context_lines is None else context_lines + logs = fetch_and_filter_logs( + ctx=ctx, + lines=num_lines, + since_time=since_time, + until_time=until_time, + ) + return search_and_merge_logs( + logs=logs, + keyword=keyword, + context_lines=ctx_lines, + is_regex=is_regex, + ) + except Exception as e: # pylint: disable=broad-exception-caught + logger.error(f"Failed to search logs: {e}") + return f"Error searching logs: {e}" + + +# Universal tool instance & aliases +search_logs = SearchLogsTool() +SearchLogs = SearchLogsTool +SearchLogsToolAlias = SearchLogsTool +ToolRegistry.register(search_logs) + + +def get_search_logs_tool(ctx: ArtemisContext) -> BaseTool: + """Exports search_logs as a LangChain BaseTool.""" + return trace_langchain_tool(search_logs.to_langchain_tool(ctx), ctx) diff --git a/artemis/tools/object_detection_tool.py b/artemis/tools/object_detection_tool.py new file mode 100644 index 0000000..234364a --- /dev/null +++ b/artemis/tools/object_detection_tool.py @@ -0,0 +1,246 @@ +# 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. + +"""Universal Object Detection tool for visual element detection and localization.""" + +import json +from pathlib import Path +from typing import Any + +from langchain_core.tools import BaseTool +from pydantic import BaseModel, Field + +from artemis.agents.object_detector.object_detector import ( + _create_error_command, + _create_success_command, + _run_object_detection, +) +from artemis.context import ArtemisContext +from artemis.data_engine.trace import ( + TraceSpan, + 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.tools.types import CyFunctionDetector +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +# pylint: disable=too-few-public-methods +class ObjectDetectionArgs(BaseModel): + """Arguments schema for object detection with explicit screenshot path.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + image_path: str = Field(..., description="Absolute path to the screenshot file.") + queries: list[str] = Field(..., description="A list of single query strings to find.") + + +# pylint: disable=too-few-public-methods +class OperatorObjectDetectionArgs(BaseModel): + """Arguments schema for operator object detection using latest screenshot from state.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + queries: list[str] = Field(..., description="A list of single query strings to find.") + + +def _load_prompt_templates() -> list[str]: + """Loads prompt templates and instructions from object_detector.json config.""" + prompt_path = Path(__file__).parent.parent.joinpath( + "agents", "object_detector", "object_detector.json" + ) + try: + with open(prompt_path, encoding="utf-8") as f: + prompt_config = json.load(f) + base_templates = prompt_config.get("templates", []) + instructions = prompt_config.get("instructions", "") + except Exception as e: # pylint: disable=broad-exception-caught + logger.warning(f"Failed to load prompt config from JSON: {e}. Using fallbacks.") + base_templates = ["Point to the following objects in the provided image: {labels_str}."] + instructions = "" + + return [f"{t}\n\n{instructions}" for t in base_templates] + + +# pylint: disable=too-many-arguments,too-many-positional-arguments,broad-exception-caught +async def _run_object_detection_logic( + ctx: ArtemisContext | None, + image_path: str | None, + queries: list[str] | None, + state: State | None = None, + tool_call_id: str | None = None, + wrapper: ToolWrapper | None = None, +) -> Any: + """Core async execution logic for object detection.""" + logger.info("Object detection tool called.") + active_wrapper = wrapper or object_detection_wrapper + templates = _load_prompt_templates() + target_image = image_path + if not target_image and state and hasattr(state, "latest_screenshot"): + target_image = state.latest_screenshot + + if not target_image: + error_msg = "No image path provided and no screenshot available in state." + if state and callable(getattr(state, "asanitize_update", None)): + return await _create_error_command(ctx, state, tool_call_id, error_msg, active_wrapper) + return f"Object Detection failed: {error_msg}" + + target_queries = queries or [] + + try: + with TraceSpan(name="run_object_detection", ctx=ctx) as span: + result = await _run_object_detection(ctx, target_image, target_queries, templates) + span.result = result + output = json.dumps(result) if not isinstance(result, str) else result + + if state and callable(getattr(state, "asanitize_update", None)): + return await _create_success_command(ctx, state, tool_call_id, output, active_wrapper) + return output + except Exception as e: + logger.error(f"Object detection execution failed: {e}") + if state and callable(getattr(state, "asanitize_update", None)): + return await _create_error_command(ctx, state, tool_call_id, str(e), active_wrapper) + return f"Object Detection failed: {e}" + + +class ObjectDetectionTool(ArtemisTool): + """Universal tool for detecting visual elements or icons in a mobile screen image.""" + + def __init__( + self, + name: str = "object_detection", + description: str | None = None, + args_schema: type[BaseModel] = ObjectDetectionArgs, + category: ToolCategory = "perception", + ): + super().__init__( + name=name, + description=description + or ( + "[EXPLORER] Detects specific visual elements or icons in a mobile screen image. " + "Use this tool when you need to find the coordinates of UI elements like " + "buttons, icons, or text fields." + ), + args_schema=args_schema, + category=category, + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + ctx: ArtemisContext | None = None, + image_path: str | None = None, + queries: list[str] | None = None, + state: State | None = None, + tool_call_id: str | None = None, + **kwargs: Any, + ) -> Any: + img_path = ( + image_path + if image_path is not None + else (kwargs.get("image_path") or kwargs.get("ImagePath")) + ) + qs = ( + queries + if queries is not None + else (kwargs.get("queries") or kwargs.get("Queries") or []) + ) + st = state if state is not None else kwargs.get("state") + tcid = tool_call_id if tool_call_id is not None else kwargs.get("tool_call_id") + wrapper = kwargs.get("wrapper") or object_detection_wrapper + return await _run_object_detection_logic( + ctx=ctx, + image_path=img_path, + queries=qs, + state=st, + tool_call_id=tcid, + wrapper=wrapper, + ) + + +class OperatorObjectDetectionTool(ObjectDetectionTool): + """Universal tool for operator detecting visual elements using screenshot from state.""" + + def __init__(self): + super().__init__( + name="object_detection", + description=( + "[EXPLORER] Detects specific visual elements or icons in a mobile screen image. " + "Use this tool when you need to find the coordinates of UI elements like " + "buttons, icons, or text fields." + ), + args_schema=OperatorObjectDetectionArgs, + category="perception", + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + ctx: ArtemisContext | None = None, + image_path: str | None = None, + queries: list[str] | None = None, + state: State | None = None, + tool_call_id: str | None = None, + **kwargs: Any, + ) -> Any: + return await super().execute( + driver=driver, + ctx=ctx, + image_path=image_path, + queries=queries, + state=state, + tool_call_id=tool_call_id, + wrapper=operator_object_detection_wrapper, + **kwargs, + ) + + +# Universal tool instances & aliases +object_detection = ObjectDetectionTool() +ObjectDetection = ObjectDetectionTool +ObjectDetectorTool = ObjectDetectionTool + +operator_object_detection = OperatorObjectDetectionTool() +OperatorObjectDetection = OperatorObjectDetectionTool +OperatorObjectDetectorTool = OperatorObjectDetectionTool + +ToolRegistry.register(object_detection) + + +def get_object_detector_tool(ctx: ArtemisContext) -> BaseTool: + """Exports object_detection as a LangChain BaseTool.""" + return trace_langchain_tool(object_detection.to_langchain_tool(ctx), ctx) + + +def get_operator_object_detector_tool(ctx: ArtemisContext) -> BaseTool: + """Exports operator_object_detection as a LangChain BaseTool.""" + return trace_langchain_tool(operator_object_detection.to_langchain_tool(ctx), ctx) + + +object_detection_wrapper = ToolWrapper( + tool_fn_getter=get_object_detector_tool, + on_success_fn=lambda output: f"Object Detection successful. Results:\n{output}", + on_failure_fn=lambda error: f"Object Detection failed: {error}", +) + +operator_object_detection_wrapper = ToolWrapper( + tool_fn_getter=get_operator_object_detector_tool, + on_success_fn=lambda output: f"Object Detection successful. Results:\n{output}", + on_failure_fn=lambda error: f"Object Detection failed: {error}", +) diff --git a/artemis/tools/perception/object_detect.py b/artemis/tools/perception/object_detect.py new file mode 100644 index 0000000..76716ea --- /dev/null +++ b/artemis/tools/perception/object_detect.py @@ -0,0 +1,43 @@ +# 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. + +"""Visual Object and Icon Detection perception tool.""" + +from pydantic import BaseModel, Field +from artemis.tools.base import artemis_tool +from artemis.drivers.base import BaseDeviceDriver + + +class ObjectDetectArgs(BaseModel): + """Arguments schema for object detection.""" + + icon_description: str = Field( + ..., description="Description of the icon, button, or graphic element." + ) + + +@artemis_tool( + name="detect_object", + description=( + "[PERCEPTION] Detects visual icon or UI element on screen using multimodal localization." + ), + args_schema=ObjectDetectArgs, + category="perception", +) +async def detect_object( + icon_description: str, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument +) -> str: + """Detects visual icon or UI element on screen using multimodal localization.""" + return f"Object Detect: Found '{icon_description}' at normalized [500, 500]" diff --git a/artemis/tools/perception/ocr.py b/artemis/tools/perception/ocr.py new file mode 100644 index 0000000..e88c7b7 --- /dev/null +++ b/artemis/tools/perception/ocr.py @@ -0,0 +1,46 @@ +# 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. + +"""Optical Character Recognition (OCR) perception tool.""" + +from pydantic import BaseModel, Field +from artemis.tools.base import artemis_tool +from artemis.drivers.base import BaseDeviceDriver + + +class OCRRecognitionArgs(BaseModel): + """Arguments schema for OCR recognition.""" + + query_text: str | None = Field( + default=None, + description="Optional target text snippet to search on screen.", + ) + + +@artemis_tool( + name="ocr_recognize", + description=( + "[PERCEPTION] Performs OCR on the active screen to extract visible text and coordinates." + ), + args_schema=OCRRecognitionArgs, + category="perception", +) +async def ocr_recognize( + query_text: str | None = None, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument +) -> str: + """Performs OCR on the active screen to extract visible text and coordinates.""" + if query_text: + return f"OCR: Text '{query_text}' located at center [500, 500]" + return "OCR: Extracted visible text items from screen." diff --git a/artemis/tools/perception/xml_search.py b/artemis/tools/perception/xml_search.py new file mode 100644 index 0000000..cb13c97 --- /dev/null +++ b/artemis/tools/perception/xml_search.py @@ -0,0 +1,45 @@ +# 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. + +"""XML Layout Search perception tool.""" + +from pydantic import BaseModel, Field +from artemis.tools.base import artemis_tool +from artemis.drivers.base import BaseDeviceDriver + + +class XMLSearchArgs(BaseModel): + """Arguments schema for XML layout search.""" + + query: str = Field( + ..., + description="Text, resource ID, or class pattern to search in UI hierarchy.", + ) + + +@artemis_tool( + name="xml_search", + description=( + "[PERCEPTION] Searches Android UI layout hierarchy for elements" + " matching text or resource ID." + ), + args_schema=XMLSearchArgs, + category="perception", +) +async def xml_search( + query: str, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument +) -> str: + """Searches Android UI layout hierarchy for elements matching text or resource ID.""" + return f"XML Search: Found element matching '{query}'" diff --git a/artemis/tools/scratchpad.py b/artemis/tools/scratchpad.py new file mode 100644 index 0000000..dcc946c --- /dev/null +++ b/artemis/tools/scratchpad.py @@ -0,0 +1,903 @@ +# 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. + +from typing import Any + +from langchain_core.messages import ToolMessage +from langchain_core.tools.base import BaseTool +from langgraph.types import Command +from pydantic import BaseModel, Field + +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, ToolRegistry +from artemis.tools.tool_wrapper import ToolWrapper +from artemis.tools.types import CyFunctionDetector +from artemis.utils.logger import get_logger +from artemis.utils.notes import ( + APPEND_NOTE_ARG_CONTENT_DESC, + APPEND_NOTE_ARG_KEY_DESC, + APPEND_NOTE_DOCSTRING, + LIST_NOTES_DOCSTRING, + READ_NOTE_ARG_KEY_DESC, + READ_NOTE_DOCSTRING, + SAVE_NOTE_ARG_CONTENT_DESC, + SAVE_NOTE_ARG_KEY_DESC, + SAVE_NOTE_DOCSTRING, + UPDATE_NOTE_ARG_KEY_DESC, + UPDATE_NOTE_ARG_REPLACEMENT_DESC, + UPDATE_NOTE_ARG_TARGET_DESC, + UPDATE_NOTE_DOCSTRING, + append_note_content, + format_list_notes_failure, + format_list_notes_success, + format_read_note_failure, + format_read_note_success, + list_notes_info, + read_note_content, + save_note_content, + update_note_content, +) + +logger = get_logger(__name__) + + +class SaveNoteArgs(BaseModel): + """Arguments schema for saving notes.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + key: str = Field(..., description=SAVE_NOTE_ARG_KEY_DESC) + content: str = Field(..., description=SAVE_NOTE_ARG_CONTENT_DESC) + + +class AppendNoteArgs(BaseModel): + """Arguments schema for appending notes.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + key: str = Field(..., description=APPEND_NOTE_ARG_KEY_DESC) + content: str = Field(..., description=APPEND_NOTE_ARG_CONTENT_DESC) + + +class UpdateNoteArgs(BaseModel): + """Arguments schema for updating notes.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + key: str = Field(..., description=UPDATE_NOTE_ARG_KEY_DESC) + target: str = Field(..., description=UPDATE_NOTE_ARG_TARGET_DESC) + replacement: str = Field(..., description=UPDATE_NOTE_ARG_REPLACEMENT_DESC) + + +class ReadNoteArgs(BaseModel): + """Arguments schema for reading notes.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + key: str = Field(..., description=READ_NOTE_ARG_KEY_DESC) + start_line: int | None = Field(None, description="Start line to read (1-indexed, inclusive)") + end_line: int | None = Field(None, description="End line to read (1-indexed, inclusive)") + + +class ListNotesArgs(BaseModel): + """Arguments schema for listing notes.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + + +class SaveNoteTool(ArtemisTool): + """Universal tool for saving notes to persistent memory.""" + + def __init__(self): + super().__init__( + name="save_note", + description=SAVE_NOTE_DOCSTRING, + args_schema=SaveNoteArgs, + category="memory", + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + ctx: ArtemisContext | None = None, + key: str | None = None, + content: str | None = None, + tool_call_id: str | None = None, + state: State | None = None, + **kwargs: Any, + ) -> Any: + k = ( + key + if key is not None + else ( + kwargs.get("key") + or kwargs.get("Key") + or kwargs.get("title") + or kwargs.get("Title") + or "" + ) + ) + c = ( + content + if content is not None + else (kwargs.get("content") or kwargs.get("Content") or "") + ) + tcid = (tool_call_id if tool_call_id is not None else kwargs.get("tool_call_id")) or "" + st = state if state is not None else kwargs.get("state") + + try: + logger.info(f"save_note called with key='{k}'") + base_dir = ( + ctx.data_engine.base_dir + if ( + ctx + and getattr(ctx, "data_engine", None) + and getattr(ctx.data_engine, "base_dir", None) + ) + else "." + ) + save_note_content(base_dir, k, c) + agent_outcome = f"Successfully saved note to {k}.md." + status = "success" + except Exception as e: # pylint: disable=broad-exception-caught + agent_outcome = f"Failed to save note {k}.md: {e}" + status = "error" + + if st and callable(getattr(st, "asanitize_update", None)): + tool_message = ToolMessage( + tool_call_id=tcid, + content=agent_outcome, + status=status, + ) + + current_agent = "operator" + if hasattr(st, "current_agent") and st.current_agent: + current_agent = st.current_agent + + return Command( + update=await st.asanitize_update( + ctx=ctx, + update={ + VALIDATOR_MESSAGES_KEY: [tool_message], + }, + agent=current_agent, + ), + ) + + return agent_outcome + + +# Universal tool instance & aliases +save_note = SaveNoteTool() +SaveNote = SaveNoteTool +ToolRegistry.register(save_note) + + +def get_save_note_tool(ctx: ArtemisContext) -> BaseTool: + """Exports save_note as a LangChain BaseTool.""" + return trace_langchain_tool(save_note.to_langchain_tool(ctx), ctx) + + +class AppendNoteTool(ArtemisTool): + """Universal tool for appending text content to notes in persistent memory.""" + + def __init__(self): + super().__init__( + name="append_note", + description=APPEND_NOTE_DOCSTRING, + args_schema=AppendNoteArgs, + category="memory", + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + ctx: ArtemisContext | None = None, + key: str | None = None, + content: str | None = None, + tool_call_id: str | None = None, + state: State | None = None, + **kwargs: Any, + ) -> Any: + k = ( + key + if key is not None + else ( + kwargs.get("key") + or kwargs.get("Key") + or kwargs.get("title") + or kwargs.get("Title") + or "" + ) + ) + c = ( + content + if content is not None + else (kwargs.get("content") or kwargs.get("Content") or "") + ) + tcid = (tool_call_id if tool_call_id is not None else kwargs.get("tool_call_id")) or "" + st = state if state is not None else kwargs.get("state") + + try: + logger.info(f"append_note called with key='{k}'") + base_dir = ( + ctx.data_engine.base_dir + if ( + ctx + and getattr(ctx, "data_engine", None) + and getattr(ctx.data_engine, "base_dir", None) + ) + else "." + ) + append_note_content(base_dir, k, c) + agent_outcome = f"Successfully appended note to {k}.md." + status = "success" + except Exception as e: # pylint: disable=broad-exception-caught + agent_outcome = f"Failed to append note {k}.md: {e}" + status = "error" + + if st and callable(getattr(st, "asanitize_update", None)): + tool_message = ToolMessage( + tool_call_id=tcid, + content=agent_outcome, + status=status, + ) + + current_agent = "operator" + if hasattr(st, "current_agent") and st.current_agent: + current_agent = st.current_agent + + return Command( + update=await st.asanitize_update( + ctx=ctx, + update={ + VALIDATOR_MESSAGES_KEY: [tool_message], + }, + agent=current_agent, + ), + ) + + return agent_outcome + + +# Universal tool instance & aliases +append_note = AppendNoteTool() +AppendNote = AppendNoteTool +ToolRegistry.register(append_note) + + +def get_append_note_tool(ctx: ArtemisContext) -> BaseTool: + """Exports append_note as a LangChain BaseTool.""" + return trace_langchain_tool(append_note.to_langchain_tool(ctx), ctx) + + +class UpdateNoteTool(ArtemisTool): + """Universal tool for updating notes by replacing a target string in persistent memory.""" + + def __init__(self): + super().__init__( + name="update_note", + description=UPDATE_NOTE_DOCSTRING, + args_schema=UpdateNoteArgs, + category="memory", + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + ctx: ArtemisContext | None = None, + key: str | None = None, + target: str | None = None, + replacement: str | None = None, + tool_call_id: str | None = None, + state: State | None = None, + **kwargs: Any, + ) -> Any: + k = ( + key + if key is not None + else ( + kwargs.get("key") + or kwargs.get("Key") + or kwargs.get("title") + or kwargs.get("Title") + or "" + ) + ) + t = target if target is not None else (kwargs.get("target") or kwargs.get("Target") or "") + r = ( + replacement + if replacement is not None + else (kwargs.get("replacement") or kwargs.get("Replacement") or "") + ) + tcid = (tool_call_id if tool_call_id is not None else kwargs.get("tool_call_id")) or "" + st = state if state is not None else kwargs.get("state") + + try: + logger.info(f"update_note called with key='{k}'") + base_dir = ( + ctx.data_engine.base_dir + if ( + ctx + and getattr(ctx, "data_engine", None) + and getattr(ctx.data_engine, "base_dir", None) + ) + else "." + ) + warning = update_note_content(base_dir, k, t, r) + if warning: + agent_outcome = f"Successfully updated note '{k}'.\nWARNING: {warning}" + else: + agent_outcome = f"Successfully updated note '{k}'." + status = "success" + except Exception as e: # pylint: disable=broad-exception-caught + agent_outcome = f"Failed to update note '{k}': {e}" + status = "error" + + if st and callable(getattr(st, "asanitize_update", None)): + tool_message = ToolMessage( + tool_call_id=tcid, + content=agent_outcome, + status=status, + ) + + current_agent = "operator" + if hasattr(st, "current_agent") and st.current_agent: + current_agent = st.current_agent + + return Command( + update=await st.asanitize_update( + ctx=ctx, + update={ + VALIDATOR_MESSAGES_KEY: [tool_message], + }, + agent=current_agent, + ), + ) + + return agent_outcome + + +# Universal tool instance & aliases +update_note = UpdateNoteTool() +UpdateNote = UpdateNoteTool +ToolRegistry.register(update_note) + + +def get_update_note_tool(ctx: ArtemisContext) -> BaseTool: + """Exports update_note as a LangChain BaseTool.""" + return trace_langchain_tool(update_note.to_langchain_tool(ctx), ctx) + + +class ReadNoteTool(ArtemisTool): + """Universal tool for reading previously saved notes from persistent memory.""" + + def __init__(self): + super().__init__( + name="read_note", + description=READ_NOTE_DOCSTRING, + args_schema=ReadNoteArgs, + category="memory", + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + ctx: ArtemisContext | None = None, + key: str | None = None, + start_line: int | None = None, + end_line: int | None = None, + tool_call_id: str | None = None, + state: State | None = None, + **kwargs: Any, + ) -> Any: + k = ( + key + if key is not None + else ( + kwargs.get("key") + or kwargs.get("Key") + or kwargs.get("title") + or kwargs.get("Title") + or "" + ) + ) + sl = ( + start_line + if start_line is not None + else (kwargs.get("start_line") or kwargs.get("StartLine")) + ) + el = end_line if end_line is not None else (kwargs.get("end_line") or kwargs.get("EndLine")) + tcid = (tool_call_id if tool_call_id is not None else kwargs.get("tool_call_id")) or "" + st = state if state is not None else kwargs.get("state") + + try: + logger.info(f"read_note called with key='{k}'") + base_dir = ( + ctx.data_engine.base_dir + if ( + ctx + and getattr(ctx, "data_engine", None) + and getattr(ctx.data_engine, "base_dir", None) + ) + else "." + ) + content = read_note_content(base_dir, k, sl, el) + agent_outcome = read_note_wrapper.on_success_fn(k, content, sl, el) + status = "success" + except Exception as e: # pylint: disable=broad-exception-caught + if "not found" in str(e).lower(): + agent_outcome = read_note_wrapper.on_failure_fn(k) + else: + agent_outcome = f"Failed to read note '{k}': {e}" + status = "error" + + if st and callable(getattr(st, "asanitize_update", None)): + tool_message = ToolMessage( + tool_call_id=tcid, + content=agent_outcome, + status=status, + ) + + current_agent = "operator" + if hasattr(st, "current_agent") and st.current_agent: + current_agent = st.current_agent + + return Command( + update=await st.asanitize_update( + ctx=ctx, + update={ + VALIDATOR_MESSAGES_KEY: [tool_message], + }, + agent=current_agent, + ), + ) + + return agent_outcome + + +# Universal tool instance & aliases +read_note = ReadNoteTool() +ReadNote = ReadNoteTool +ToolRegistry.register(read_note) + + +def get_read_note_tool(ctx: ArtemisContext) -> BaseTool: + """Exports read_note as a LangChain BaseTool.""" + return trace_langchain_tool(read_note.to_langchain_tool(ctx), ctx) + + +class ListNotesTool(ArtemisTool): + """Universal tool for listing all stored note keys in persistent memory.""" + + def __init__(self): + super().__init__( + name="list_notes", + description=LIST_NOTES_DOCSTRING, + args_schema=ListNotesArgs, + category="memory", + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + 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")) or "" + st = state if state is not None else kwargs.get("state") + + try: + logger.info("list_notes called") + base_dir = ( + ctx.data_engine.base_dir + if ( + ctx + and getattr(ctx, "data_engine", None) + and getattr(ctx.data_engine, "base_dir", None) + ) + else "." + ) + notes_info = list_notes_info(base_dir) + agent_outcome = list_notes_wrapper.on_success_fn(notes_info) + status = "success" + except Exception as e: # pylint: disable=broad-exception-caught + agent_outcome = f"Failed to list notes: {e}" + status = "error" + + if st and callable(getattr(st, "asanitize_update", None)): + tool_message = ToolMessage( + tool_call_id=tcid, + content=agent_outcome, + status=status, + ) + + current_agent = "operator" + if hasattr(st, "current_agent") and st.current_agent: + current_agent = st.current_agent + + return Command( + update=await st.asanitize_update( + ctx=ctx, + update={ + VALIDATOR_MESSAGES_KEY: [tool_message], + }, + agent=current_agent, + ), + ) + + return agent_outcome + + +# Universal tool instance & aliases +list_notes = ListNotesTool() +ListNotes = ListNotesTool +ToolRegistry.register(list_notes) + + +def get_list_notes_tool(ctx: ArtemisContext) -> BaseTool: + """Exports list_notes as a LangChain BaseTool.""" + return trace_langchain_tool(list_notes.to_langchain_tool(ctx), ctx) + + +class ReadNotePureTool(ArtemisTool): + """Universal pure tool for reading previously saved notes from persistent memory.""" + + def __init__(self): + super().__init__( + name="read_note_pure", + description=READ_NOTE_DOCSTRING, + args_schema=ReadNoteArgs, + category="memory", + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + ctx: ArtemisContext | None = None, + key: str | None = None, + start_line: int | None = None, + end_line: int | None = None, + **kwargs: Any, + ) -> str: + k = ( + key + if key is not None + else ( + kwargs.get("key") + or kwargs.get("Key") + or kwargs.get("title") + or kwargs.get("Title") + or "" + ) + ) + sl = ( + start_line + if start_line is not None + else (kwargs.get("start_line") or kwargs.get("StartLine")) + ) + el = end_line if end_line is not None else (kwargs.get("end_line") or kwargs.get("EndLine")) + + try: + logger.info(f"read_note_pure called with key='{k}'") + base_dir = ( + ctx.data_engine.base_dir + if ( + ctx + and getattr(ctx, "data_engine", None) + and getattr(ctx.data_engine, "base_dir", None) + ) + else "." + ) + content = read_note_content(base_dir, k, sl, el) + return format_read_note_success(k, content, sl, el) + except Exception as e: # pylint: disable=broad-exception-caught + return format_read_note_failure(k, str(e)) + + +# Universal pure tool instance & aliases +read_note_pure = ReadNotePureTool() +ReadNotePure = ReadNotePureTool +ToolRegistry.register(read_note_pure) + + +def get_read_note_tool_pure(ctx: ArtemisContext) -> BaseTool: + """Exports read_note_pure as a LangChain BaseTool named 'read_note'.""" + return read_note_pure.to_langchain_tool(ctx, name="read_note") + + +class ListNotesPureTool(ArtemisTool): + """Universal pure tool for listing all stored note keys in persistent memory.""" + + def __init__(self): + super().__init__( + name="list_notes_pure", + description=LIST_NOTES_DOCSTRING, + args_schema=ListNotesArgs, + category="memory", + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + ctx: ArtemisContext | None = None, + **kwargs: Any, + ) -> str: + try: + logger.info("list_notes_pure called") + base_dir = ( + ctx.data_engine.base_dir + if ( + ctx + and getattr(ctx, "data_engine", None) + and getattr(ctx.data_engine, "base_dir", None) + ) + else "." + ) + notes_info = list_notes_info(base_dir) + return format_list_notes_success(notes_info) + except Exception as e: # pylint: disable=broad-exception-caught + return format_list_notes_failure(str(e)) + + +# Universal pure tool instance & aliases +list_notes_pure = ListNotesPureTool() +ListNotesPure = ListNotesPureTool +ToolRegistry.register(list_notes_pure) + + +def get_list_notes_tool_pure(ctx: ArtemisContext) -> BaseTool: + """Exports list_notes_pure as a LangChain BaseTool named 'list_notes'.""" + return list_notes_pure.to_langchain_tool(ctx, name="list_notes") + + +class SaveNotePureTool(ArtemisTool): + """Universal pure tool for saving notes to persistent memory.""" + + def __init__(self): + super().__init__( + name="save_note_pure", + description=SAVE_NOTE_DOCSTRING, + args_schema=SaveNoteArgs, + category="memory", + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + ctx: ArtemisContext | None = None, + key: str | None = None, + content: str | None = None, + **kwargs: Any, + ) -> str: + k = ( + key + if key is not None + else ( + kwargs.get("key") + or kwargs.get("Key") + or kwargs.get("title") + or kwargs.get("Title") + or "" + ) + ) + c = ( + content + if content is not None + else (kwargs.get("content") or kwargs.get("Content") or "") + ) + + try: + logger.info(f"save_note_pure called with key='{k}'") + base_dir = ( + ctx.data_engine.base_dir + if ( + ctx + and getattr(ctx, "data_engine", None) + and getattr(ctx.data_engine, "base_dir", None) + ) + else "." + ) + save_note_content(base_dir, k, c) + return f"Successfully saved note to {k}.md." + except Exception as e: # pylint: disable=broad-exception-caught + return f"Failed to save note {k}.md: {e}" + + +# Universal pure tool instance & aliases +save_note_pure = SaveNotePureTool() +SaveNotePure = SaveNotePureTool +ToolRegistry.register(save_note_pure) + + +def get_save_note_tool_pure(ctx: ArtemisContext) -> BaseTool: + """Exports save_note_pure as a LangChain BaseTool named 'save_note'.""" + return save_note_pure.to_langchain_tool(ctx, name="save_note") + + +class UpdateNotePureTool(ArtemisTool): + """Universal pure tool for updating notes by replacing a target string in persistent memory.""" + + def __init__(self): + super().__init__( + name="update_note_pure", + description=UPDATE_NOTE_DOCSTRING, + args_schema=UpdateNoteArgs, + category="memory", + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + ctx: ArtemisContext | None = None, + key: str | None = None, + target: str | None = None, + replacement: str | None = None, + **kwargs: Any, + ) -> str: + k = ( + key + if key is not None + else ( + kwargs.get("key") + or kwargs.get("Key") + or kwargs.get("title") + or kwargs.get("Title") + or "" + ) + ) + t = target if target is not None else (kwargs.get("target") or kwargs.get("Target") or "") + r = ( + replacement + if replacement is not None + else (kwargs.get("replacement") or kwargs.get("Replacement") or "") + ) + + try: + logger.info(f"update_note_pure called with key='{k}'") + base_dir = ( + ctx.data_engine.base_dir + if ( + ctx + and getattr(ctx, "data_engine", None) + and getattr(ctx.data_engine, "base_dir", None) + ) + else "." + ) + warning = update_note_content(base_dir, k, t, r) + if warning: + return f"Successfully updated note '{k}'.\nWARNING: {warning}" + return f"Successfully updated note '{k}'." + except Exception as e: # pylint: disable=broad-exception-caught + return f"Failed to update note '{k}': {e}" + + +# Universal pure tool instance & aliases +update_note_pure = UpdateNotePureTool() +UpdateNotePure = UpdateNotePureTool +ToolRegistry.register(update_note_pure) + + +def get_update_note_tool_pure(ctx: ArtemisContext) -> BaseTool: + """Exports update_note_pure as a LangChain BaseTool named 'update_note'.""" + return update_note_pure.to_langchain_tool(ctx, name="update_note") + + +class AppendNotePureTool(ArtemisTool): + """Universal pure tool for appending text content to notes in persistent memory.""" + + def __init__(self): + super().__init__( + name="append_note_pure", + description=APPEND_NOTE_DOCSTRING, + args_schema=AppendNoteArgs, + category="memory", + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + ctx: ArtemisContext | None = None, + key: str | None = None, + content: str | None = None, + **kwargs: Any, + ) -> str: + k = ( + key + if key is not None + else ( + kwargs.get("key") + or kwargs.get("Key") + or kwargs.get("title") + or kwargs.get("Title") + or "" + ) + ) + c = ( + content + if content is not None + else (kwargs.get("content") or kwargs.get("Content") or "") + ) + + try: + logger.info(f"append_note_pure called with key='{k}'") + base_dir = ( + ctx.data_engine.base_dir + if ( + ctx + and getattr(ctx, "data_engine", None) + and getattr(ctx.data_engine, "base_dir", None) + ) + else "." + ) + append_note_content(base_dir, k, c) + return f"Successfully appended note to {k}.md." + except Exception as e: # pylint: disable=broad-exception-caught + return f"Failed to append note {k}.md: {e}" + + +# Universal pure tool instance & aliases +append_note_pure = AppendNotePureTool() +AppendNotePure = AppendNotePureTool +ToolRegistry.register(append_note_pure) + + +def get_append_note_tool_pure(ctx: ArtemisContext) -> BaseTool: + """Exports append_note_pure as a LangChain BaseTool named 'append_note'.""" + return append_note_pure.to_langchain_tool(ctx, name="append_note") + + +save_note_wrapper = ToolWrapper( + tool_fn_getter=get_save_note_tool, + on_success_fn=lambda key: f"Successfully saved note '{key}'.", + on_failure_fn=lambda key: f"Failed to save note '{key}'.", +) + +read_note_wrapper = ToolWrapper( + tool_fn_getter=get_read_note_tool, + on_success_fn=format_read_note_success, + on_failure_fn=lambda key: format_read_note_failure(key, "not found"), +) + +list_notes_wrapper = ToolWrapper( + tool_fn_getter=get_list_notes_tool, + on_success_fn=format_list_notes_success, + on_failure_fn=lambda: format_list_notes_failure("error"), +) + + +update_note_wrapper = ToolWrapper( + tool_fn_getter=get_update_note_tool, + on_success_fn=lambda key: f"Successfully updated note '{key}'.", + on_failure_fn=lambda key: f"Failed to update note '{key}'.", +) + + +append_note_wrapper = ToolWrapper( + tool_fn_getter=get_append_note_tool, + on_success_fn=lambda key: f"Successfully appended note '{key}'.", + on_failure_fn=lambda key: f"Failed to append note '{key}'.", +) diff --git a/artemis/tools/system/adb_shell.py b/artemis/tools/system/adb_shell.py new file mode 100644 index 0000000..59faa0c --- /dev/null +++ b/artemis/tools/system/adb_shell.py @@ -0,0 +1,31 @@ +# 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. + +from artemis.tools.command_tool import ( + RunAdbCommandArgs, + RunAdbCommandTool, + RunAdbCommand, + run_adb_command, +) + +# Backward-compatibility alias +AdbShellArgs = RunAdbCommandArgs + +__all__ = [ + "RunAdbCommandArgs", + "RunAdbCommandTool", + "RunAdbCommand", + "run_adb_command", + "AdbShellArgs", +] diff --git a/artemis/tools/tool_wrapper.py b/artemis/tools/tool_wrapper.py new file mode 100644 index 0000000..31c1291 --- /dev/null +++ b/artemis/tools/tool_wrapper.py @@ -0,0 +1,179 @@ +# 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. + +from collections.abc import Callable +import inspect +import time +from typing import Annotated, Any, get_args, get_origin +from unittest.mock import MagicMock, Mock +import uuid + +from langchain_core.tools import BaseTool +from langchain_core.tools.base import InjectedToolCallId +from langgraph.prebuilt import InjectedState +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 import engine as engine_mod +from artemis.data_engine.trace import CURRENT_TRACE_ID, smart_serialize +from artemis.graph.state import State +from artemis.tools.types import CyFunctionDetector + + +class ToolWrapper(BaseModel): + """Wrapper holding a tool factory and lifecycle callbacks.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + tool_fn_getter: Callable[[ArtemisContext], BaseTool] + on_success_fn: Callable[..., str] + on_failure_fn: Callable[..., str] + + +class CompositeToolWrapper(ToolWrapper): + """Wrapper holding a composite tool factory and lifecycle callbacks.""" + + composite_tools_fn_getter: Callable[[ArtemisContext], list[BaseTool]] + + +# pylint: disable=too-many-branches,too-many-locals,too-many-statements +async def invoke_tool_with_injection( + tool: BaseTool, + args: dict, + tool_call_id: str, + state: State | None = None, + record_trace: bool = True, +) -> Any: + """Invokes a LangChain tool, manually injecting arguments marked with + + InjectedState or InjectedToolCallId if calling the underlying function + directly. + Automatically records type='tool' trace in DataEngine. + """ + final_args = dict(args) + + if isinstance(tool, (Mock, MagicMock)): + has_real_coro = hasattr(tool, "coroutine") and not isinstance( + tool.coroutine, (Mock, MagicMock) + ) + has_real_func = hasattr(tool, "func") and not isinstance(tool.func, (Mock, MagicMock)) + if not has_real_coro and not has_real_func: + if ( + state is not None + and hasattr(tool, "args") + and isinstance(tool.args, dict) + and "state" in tool.args + ): + final_args["state"] = state + return await tool.ainvoke(final_args) + + # For StructuredTool, the original function might be in tool.func or tool.coroutine + func = None + if hasattr(tool, "coroutine") and tool.coroutine: + func = tool.coroutine + elif hasattr(tool, "func") and tool.func: + func = tool.func + + if func: + try: + sig = inspect.signature(func) + for param_name, param in sig.parameters.items(): + annotation = param.annotation + # Handle Annotated types + if get_origin(annotation) is Annotated: + args_list = get_args(annotation) + if InjectedState in args_list: + final_args[param_name] = state + elif InjectedToolCallId in args_list: + final_args[param_name] = tool_call_id + except Exception: # pylint: disable=broad-exception-caught + # If inspection fails, fall back to just using provided args + pass + + trace_id = uuid.uuid4() + parent_id = CURRENT_TRACE_ID.get() + token = CURRENT_TRACE_ID.set(trace_id) + start_time = time.time() + + live_engine = getattr(engine_mod, "_CURRENT_DATA_ENGINE", None) + + if record_trace and live_engine: + step_id = getattr(live_engine, "current_step_id", None) + live_engine.record_trace( + type="tool", + name=tool.name, + payload={"args": {k: smart_serialize(v) for k, v in args.items()}}, + status="running", + parent_trace_id=parent_id, + step_id=step_id, + trace_id=trace_id, + ) + + execution_status = "running" + execution_result = None + execution_error = None + try: + # Call the function + if inspect.iscoroutinefunction(func): + result = await func(**final_args) + else: + result = func(**final_args) + execution_status = "success" + execution_result = result + return result + except Exception as e: + execution_status = "failed" + execution_error = e + raise e + finally: + duration = time.time() - start_time + if record_trace and live_engine: + step_id = getattr(live_engine, "current_step_id", None) + payload = {"args": {k: smart_serialize(v) for k, v in args.items()}} + + if execution_status == "success": + payload["result"] = smart_serialize(get_tool_result_content(execution_result)) + elif execution_status == "failed": + payload["error"] = str(execution_error) + else: + execution_status = "terminated" + payload["error"] = "Execution terminated unexpectedly." + + live_engine.record_trace( + type="tool", + name=tool.name, + payload=payload, + status=execution_status, + duration=duration, + parent_trace_id=parent_id, + step_id=step_id, + trace_id=trace_id, + ) + CURRENT_TRACE_ID.reset(token) + else: + # Fallback to ainvoke if no direct func available + return await tool.ainvoke(final_args) + + +def get_tool_result_content(result: Any) -> Any: + """Extracts the content from a tool result, handling Command objects.""" + if isinstance(result, Command): + updates = result.update + if VALIDATOR_MESSAGES_KEY in updates: + msgs = updates[VALIDATOR_MESSAGES_KEY] + if msgs and hasattr(msgs[0], "content"): + return msgs[0].content + return result + return result diff --git a/artemis/tools/types.py b/artemis/tools/types.py new file mode 100644 index 0000000..90293db --- /dev/null +++ b/artemis/tools/types.py @@ -0,0 +1,72 @@ +# 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. + +from pydantic import BaseModel, Field, model_validator + +from artemis.utils.ui_hierarchy import ElementBounds + + +class _CyFunctionDetectorMeta(type): + """Metaclass for detecting cython compiled functions.""" + + def __instancecheck__(cls, instance): + name = type(instance).__name__ + return ( + name + in ( + "cyfunction", + "cython_function_or_method", + "builtin_function_or_method", + ) + or "cyfunction" in name.lower() + ) + + +# pylint: disable=too-few-public-methods +class CyFunctionDetector(metaclass=_CyFunctionDetectorMeta): + """Detector for cython compiled functions.""" + + +class Target(BaseModel): + """A comprehensive locator for a UI element, supporting a fallback mechanism.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + + resource_id: str | None = Field(None, description="The resource-id of the element.") + resource_id_index: int | None = Field( + None, + description=("The zero-based index if multiple elements share the same resource-id."), + ) + text: str | None = Field( + None, + description=("The text content of the element (e.g., a label or placeholder)."), + ) + text_index: int | None = Field( + None, + description=("The zero-based index if multiple elements share the same text."), + ) + bounds: ElementBounds | None = Field( + None, description="The x, y, width, and height of the element." + ) + + @model_validator(mode="after") + def _default_indices(self): + # Treat empty strings like "not provided" + if ( + self.resource_id is not None and self.resource_id != "" + ) and self.resource_id_index is None: + self.resource_id_index = 0 + if (self.text is not None and self.text != "") and self.text_index is None: + self.text_index = 0 + return self diff --git a/artemis/tools/utils.py b/artemis/tools/utils.py new file mode 100644 index 0000000..47e4f68 --- /dev/null +++ b/artemis/tools/utils.py @@ -0,0 +1,389 @@ +# 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. + +from __future__ import annotations + +from typing import Literal + +from pydantic import BaseModel, ConfigDict + +from artemis.context import ArtemisContext +from artemis.controllers.types import ( + CoordinatesSelectorRequest, + PercentagesSelectorRequest, +) +from artemis.controllers.unified_controller import UnifiedMobileController +from artemis.graph.state import State +from artemis.tools.types import Target +from artemis.utils.logger import get_logger +from artemis.utils.ui_hierarchy import ( + ElementBounds, + Point, + find_element_by_resource_id, + get_bounds_for_element, + get_element_text, + is_element_focused, +) + +logger = get_logger(__name__) + + +def find_element_by_text( + ui_hierarchy: list[dict], text: str, index: int | None = None +) -> dict | None: + """Find a UI element by its text content (adapted to both flat and rich hierarchy) + + This function performs a recursive, case-insensitive search. + + Args: + ui_hierarchy: List of UI element dictionaries. + text: The text content to search for. + index: Optional index to select nth matching element. + + Returns: + The complete UI element dictionary if found, None otherwise. + """ + matches: list[dict] = [] + + def search_recursive(elements: list[dict]) -> None: + for element in elements: + if isinstance(element, dict): + src = element.get("attributes", element) + element_text = src.get("text", "") + + # Guard against non-string values + if not isinstance(element_text, str): + element_text = "" + + text_lower = text.lower() + if element_text and text_lower == element_text.lower(): + matches.append(element) + + if children := element.get("children", []): + search_recursive(children) + + search_recursive(ui_hierarchy) + + if not matches: + return None + + if index is None: + idx = 0 + elif index < 0: + return None + else: + idx = index + if idx < len(matches): + return matches[idx] + return None + + +async def tap_bottom_right_of_element(bounds: ElementBounds, ctx: ArtemisContext): + """Tap the bottom-right corner of the specified element bounds.""" + bottom_right: Point = bounds.get_relative_point(x_percent=0.99, y_percent=0.99) + await tap( + ctx=ctx, + selector_request=SelectorRequestWithCoordinates( + coordinates=CoordinatesSelectorRequest( + x=bottom_right.x, + y=bottom_right.y, + ), + ), + ) + + +# pylint: disable=too-many-return-statements +async def move_cursor_to_end_if_bounds( + ctx: ArtemisContext, + state: State, + target: Target, + elt: dict | None = None, +) -> dict | None: + """Best-effort move of the text cursor near the end of the input by tapping the + + bottom-right area of the focused element (if bounds are available). + """ + if target.resource_id: + if not elt: + elt = find_element_by_resource_id( + ui_hierarchy=state.latest_ui_hierarchy or [], + resource_id=target.resource_id, + index=target.resource_id_index, + ) + if not elt: + return None + + bounds = get_bounds_for_element(elt) + if not bounds: + return elt + + logger.debug("Tapping near the end of the input to move the cursor") + await tap_bottom_right_of_element(bounds=bounds, ctx=ctx) + logger.debug(f"Tapped end of input {target.resource_id}") + return elt + + if target.bounds: + await tap_bottom_right_of_element(target.bounds, ctx=ctx) + logger.debug("Tapped end of input by coordinates") + return elt + + if target.text: + text_elt = find_element_by_text( + state.latest_ui_hierarchy or [], + target.text, + index=target.text_index, + ) + if text_elt: + bounds = get_bounds_for_element(text_elt) + if bounds: + await tap_bottom_right_of_element(bounds=bounds, ctx=ctx) + logger.debug(f"Tapped end of input that had text'{target.text}'") + return text_elt + return None + + return None + + +async def focus_element_if_needed( + ctx: ArtemisContext, target: Target +) -> Literal["resource_id", "coordinates", "text", "already_focused"] | None: + """Ensures the element is focused, with a sanity check to prevent trusting misleading IDs. + + If no target locator is provided (empty target), assumes the field is + already focused + and returns "already_focused" to allow typing without re-focusing. + """ + # If no locator is provided, assume already focused (e.g., keyboard is visible) + if not target.resource_id and not target.bounds and not target.text: + logger.debug("No target locator provided, assuming element is already focused") + return "already_focused" + + controller = UnifiedMobileController(ctx) + rich_hierarchy = await controller.get_ui_elements() + elt_from_id = None + if target.resource_id: + elt_from_id = find_element_by_resource_id( + ui_hierarchy=rich_hierarchy, + resource_id=target.resource_id, + index=target.resource_id_index, + is_rich_hierarchy=False, + ) + + if elt_from_id and target.text: + text_from_id_elt = get_element_text(elt_from_id) + if not text_from_id_elt or target.text.lower() != text_from_id_elt.lower(): + logger.warning( + f"ID '{target.resource_id}' and text '{target.text}' seem to be" + " on different elements. Ignoring the resource_id and falling" + " back to other locators." + ) + elt_from_id = None + + if elt_from_id: + if not is_element_focused(elt_from_id): + await tap( + ctx=ctx, + selector_request=IdSelectorRequest(id=target.resource_id), # type: ignore + index=target.resource_id_index, + ) + logger.debug(f"Focused (tap) on resource_id={target.resource_id}") + rich_hierarchy = await controller.get_ui_elements() + elt_from_id = find_element_by_resource_id( + ui_hierarchy=rich_hierarchy, + resource_id=target.resource_id, # type: ignore + index=target.resource_id_index, + is_rich_hierarchy=False, + ) + if elt_from_id and is_element_focused(elt_from_id): + logger.debug(f"Text input is focused: {target.resource_id}") + return "resource_id" + logger.warning(f"Failed to focus using resource_id='{target.resource_id}'. Fallback...") + + if target.bounds: + relative_point = target.bounds.get_center() + await tap( + ctx=ctx, + selector_request=SelectorRequestWithCoordinates( + coordinates=CoordinatesSelectorRequest(x=relative_point.x, y=relative_point.y) + ), + ) + logger.debug(f"Tapped on coordinates ({relative_point.x}, {relative_point.y}) to focus.") + return "coordinates" + + if target.text: + text_elt = find_element_by_text(rich_hierarchy, target.text, index=target.text_index) + if text_elt: + bounds = get_bounds_for_element(text_elt) + if bounds: + relative_point = bounds.get_center() + await tap( + ctx=ctx, + selector_request=SelectorRequestWithCoordinates( + coordinates=CoordinatesSelectorRequest( + x=relative_point.x, y=relative_point.y + ) + ), + ) + logger.debug(f"Tapped on text element '{target.text}' to focus.") + return "text" + + logger.error( + "Failed to focus element. No valid locator (resource_id, coordinates, or text) succeeded." + ) + return None + + +def validate_coordinates_bounds( + target: Target, screen_width: int, screen_height: int +) -> str | None: + """Validate that coordinates are within screen bounds. + + Returns error message if invalid, None if valid. + """ + if not target.bounds: + return None + + center = target.bounds.get_center() + errors = [] + + if center.x < 0 or center.x >= screen_width: + errors.append(f"x={center.x} is outside screen width (0-{screen_width})") + if center.y < 0 or center.y >= screen_height: + errors.append(f"y={center.y} is outside screen height (0-{screen_height})") + + return "; ".join(errors) if errors else None + + +def has_valid_selectors(target: Target) -> bool: + """Check if target has at least one valid selector.""" + has_coordinates = target.bounds is not None + has_resource_id = target.resource_id is not None and target.resource_id != "" + has_text = target.text is not None and target.text != "" + return has_coordinates or has_resource_id or has_text + + +class IdSelectorRequest(BaseModel): + """Selector request using resource id.""" + + model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True) + id: str + + def to_dict(self) -> dict[str, str | int]: + """Convert selector to dictionary.""" + return {"id": self.id} + + +class TextSelectorRequest(BaseModel): + """Selector request using element text.""" + + model_config = ConfigDict(extra="forbid") + text: str + + def to_dict(self) -> dict[str, str | int]: + """Convert selector to dictionary.""" + return {"text": self.text} + + +class SelectorRequestWithCoordinates(BaseModel): + """Selector request using absolute coordinates.""" + + model_config = ConfigDict(extra="forbid") + coordinates: CoordinatesSelectorRequest + + def to_dict(self) -> dict[str, str | int]: + """Convert selector to dictionary.""" + return {"point": self.coordinates.to_str()} + + +class SelectorRequestWithPercentages(BaseModel): + """Selector request using percentage coordinates.""" + + model_config = ConfigDict(extra="forbid") + percentages: PercentagesSelectorRequest + + def to_dict(self) -> dict[str, str | int]: + """Convert selector to dictionary.""" + return {"point": self.percentages.to_str()} + + +class IdWithTextSelectorRequest(BaseModel): + """Selector request using both resource id and text.""" + + model_config = ConfigDict(extra="forbid") + id: str + text: str + + def to_dict(self) -> dict[str, str | int]: + """Convert selector to dictionary.""" + return {"id": self.id, "text": self.text} + + +SelectorRequest = ( + IdSelectorRequest + | SelectorRequestWithCoordinates + | SelectorRequestWithPercentages + | TextSelectorRequest + | IdWithTextSelectorRequest +) + + +def _extract_resource_id_and_text_from_selector( + selector: SelectorRequest, +) -> tuple[str | None, str | None]: + """Extract resource_id and text from a selector.""" + resource_id = None + text = None + + if isinstance(selector, IdSelectorRequest): + resource_id = selector.id + elif isinstance(selector, TextSelectorRequest): + text = selector.text + elif isinstance(selector, IdWithTextSelectorRequest): + resource_id = selector.id + text = selector.text + + return resource_id, text + + +async def tap( + ctx: ArtemisContext, + selector_request: SelectorRequest, + index: int | None = None, +): + """Tap on a selector. + + Index is optional and is used when you have multiple views matching the same + selector. ui_hierarchy is optional and used for ADB taps to find elements. + """ + controller = UnifiedMobileController(ctx) + if isinstance(selector_request, SelectorRequestWithCoordinates): + result = await controller.tap_at( + x=selector_request.coordinates.x, y=selector_request.coordinates.y + ) + return result.error if result.error else None + if isinstance(selector_request, SelectorRequestWithPercentages): + coords = selector_request.percentages.to_coords( + width=ctx.device.device_width, + height=ctx.device.device_height, + ) + return await controller.tap_at(coords.x, coords.y) + + # For other selectors, we need the UI hierarchy + resource_id, text = _extract_resource_id_and_text_from_selector(selector_request) + + return await controller.tap_element( + resource_id=resource_id, + text=text, + index=index if index is not None else 0, + ) diff --git a/artemis/tools/video_tool.py b/artemis/tools/video_tool.py new file mode 100644 index 0000000..d054b6e --- /dev/null +++ b/artemis/tools/video_tool.py @@ -0,0 +1,253 @@ +# 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. + +"""Video Analyzer utility for analyzing video content using Gemini models. + +This utility sends video files to video-capable Gemini models for analysis +and returns text descriptions based on the provided prompt. +It also exposes a tool for agents to perform structured video analysis. +""" + +from typing import Any, Literal + +from langchain_core.messages import ToolMessage +from langchain_core.tools import BaseTool +from langgraph.types import Command +from pydantic import BaseModel, Field + +from artemis.agents.video_analyzer.video_analyzer import VideoAnalyzer +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.tools.types import CyFunctionDetector +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +class VideoAnalyzerArgs(BaseModel): + """Arguments schema for video analyzer tool.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + time_description: str = Field( + ..., + description=( + "The specific time range to analyze, relative to the start of the" + " recording (e.g., 'from 5s to 12s', 'from 20s to the end')." + ), + ) + purpose: str = Field( + ..., + description=( + "The specific query, intent, or question you want the video" + " analyzer to answer about the visual or audio content in this" + " segment." + ), + ) + + +OPERATOR_VIDEO_ANALYZER_DOCSTRING = ( + "[VIDEO] Use this tool to analyze any specific time range of the video" + " recording. This tool delegates to a powerful video-analyzing subagent" + " with full access to the continuous screen recording of your entire" + " execution session.\n\n" + "- This tool overcomes your limitation of only receiving discrete" + " screenshots. Common use cases include (but are not limited to):\n" + " * Watching media files: Play the video on the device screen, use" + " `wait_for_delay` until the video finishes playing, and then call this" + " tool to analyze the recording.\n" + " * Bridging the gap: Inspecting what happened between your last action" + " and the current screenshot.\n" + " * Memory retrieval: Recalling fast-moving or transient information from" + " a previous step.\n" + " * Catching rapid changes: Investigating on-screen information that" + " changes too quickly for static screenshots to capture.\n" + " - Please specify a narrow time range when possible to ensure faster" + " processing.\n" + " - If only audio analysis is needed, state it clearly in your purpose.\n" +) + +DIAGNOSER_VIDEO_ANALYZER_DOCSTRING = ( + "[DIAGNOSTIC] Analyzes a video segment for failure diagnosis and" + " forensics.\n\n" + "Use this to verify if an action succeeded, check for error popups, or" + " investigate visual anomalies when a step fails.\n" + "Provide a narrow time range (e.g., 'from 10s to 15s') to ensure speed.\n" +) + + +class VideoAnalyzerTool(ArtemisTool): + """Universal tool for analyzing video content using Gemini models.""" + + def __init__( + self, + role: Literal["operator", "diagnoser"] = "operator", + description: str | None = None, + category: ToolCategory = "perception", + ): + chosen_description = description or ( + OPERATOR_VIDEO_ANALYZER_DOCSTRING + if role == "operator" + else DIAGNOSER_VIDEO_ANALYZER_DOCSTRING + ) + super().__init__( + name="video_analyzer", + description=chosen_description, + args_schema=VideoAnalyzerArgs, + category=category, + ) + self.role = role + + # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + ctx: ArtemisContext | None = None, + time_description: str | None = None, + purpose: str | None = None, + tool_call_id: str | None = None, + state: State | None = None, + **kwargs: Any, + ) -> Any: + time_desc = ( + time_description + if time_description is not None + else (kwargs.get("time_description") or kwargs.get("TimeDescription") or "") + ) + purp = ( + purpose + if purpose is not None + else (kwargs.get("purpose") or kwargs.get("Purpose") or "") + ) + 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") + + try: + logger.info(f"video_analyzer called for range: {time_desc}") + if ctx is None: + raise ValueError("ArtemisContext is required for VideoAnalyzer.") + agent = VideoAnalyzer(ctx) + agent_outcome, status = await agent.run(time_desc, purp) + except Exception as e: # pylint: disable=broad-exception-caught + logger.error(f"Error running video analyzer: {e}") + agent_outcome = f"Error running video analyzer: {e}" + status = "error" + + if st and callable(getattr(st, "asanitize_update", None)): + tool_message = ToolMessage( + tool_call_id=tcid or "", + content=agent_outcome, + status=status, + ) + return Command( + update=await st.asanitize_update( + ctx=ctx, + update={ + VALIDATOR_MESSAGES_KEY: [tool_message], + }, + agent="validator", + ), + ) + + if status == "failed": + return f"Video analysis failed: {agent_outcome}" + return agent_outcome + + +# Universal tool instance & aliases +video_analyzer = VideoAnalyzerTool() +ToolRegistry.register(video_analyzer) + + +class VideoAnalyzerPureTool(ArtemisTool): + """Universal pure tool for video analysis returning text results directly.""" + + def __init__(self): + super().__init__( + name="video_analyzer_pure", + description=OPERATOR_VIDEO_ANALYZER_DOCSTRING, + args_schema=VideoAnalyzerArgs, + category="perception", + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + ctx: ArtemisContext | None = None, + time_description: str | None = None, + purpose: str | None = None, + **kwargs: Any, + ) -> str: + time_desc = ( + time_description + if time_description is not None + else (kwargs.get("time_description") or kwargs.get("TimeDescription") or "") + ) + purp = ( + purpose + if purpose is not None + else (kwargs.get("purpose") or kwargs.get("Purpose") or "") + ) + + try: + logger.info(f"video_analyzer (pure) called for range: {time_desc}") + if ctx is None: + raise ValueError("ArtemisContext is required for VideoAnalyzer.") + agent = VideoAnalyzer(ctx) + agent_outcome, status = await agent.run(time_desc, purp) + if status == "failed": + return f"Video analysis failed: {agent_outcome}" + return agent_outcome + except Exception as e: # pylint: disable=broad-exception-caught + logger.error(f"Error in pure video_analyzer tool: {e}") + return f"Error running video analyzer: {e}" + + +# Universal pure tool instance & aliases +video_analyzer_pure = VideoAnalyzerPureTool() +VideoAnalyzerPure = VideoAnalyzerPureTool +ToolRegistry.register(video_analyzer_pure) + + +def get_video_analyzer_tool( + ctx: ArtemisContext, role: Literal["operator", "diagnoser"] = "operator" +) -> BaseTool: + """Creates a tool that allows agents to perform structured video analysis. + + The tool acts as a Main Agent following the Dynamic Video Analyzer plan. + """ + tool_inst = VideoAnalyzerTool(role=role) + return trace_langchain_tool(tool_inst.to_langchain_tool(ctx), ctx) + + +def get_video_analyzer_tool_pure(ctx: ArtemisContext) -> BaseTool: + """Creates a pure LangChain tool for video analysis that returns the text result directly. + + Suitable for agents running outside of LangGraph (like the outputter). + """ + return trace_langchain_tool( + video_analyzer_pure.to_langchain_tool(ctx, name="video_analyzer"), ctx + ) + + +video_analyzer_wrapper = ToolWrapper( + tool_fn_getter=get_video_analyzer_tool, + on_success_fn=lambda output: f"Video Analyzer replied:\n{output}", + on_failure_fn=lambda error: f"Video Analyzer failed: {error}", +) diff --git a/artemis/tools/wait_tool.py b/artemis/tools/wait_tool.py new file mode 100644 index 0000000..d3662da --- /dev/null +++ b/artemis/tools/wait_tool.py @@ -0,0 +1,98 @@ +# 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 asyncio +from typing import Any + +from langchain_core.tools import BaseTool +from pydantic import BaseModel, Field + +from artemis.context import ArtemisContext +from artemis.data_engine.trace import trace_langchain_tool +from artemis.drivers.base import BaseDeviceDriver +from artemis.tools.base import ArtemisTool, ToolCategory, ToolRegistry +from artemis.tools.tool_wrapper import ToolWrapper +from artemis.tools.types import CyFunctionDetector +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +class WaitArgs(BaseModel): + """Arguments schema for wait tool.""" + + model_config = {"ignored_types": (CyFunctionDetector,)} + seconds: int = Field( + ..., + description=("The duration to wait/sleep in seconds. Minimum is 1, maximum is 60."), + ) + + +WAIT_DOCSTRING = ( + "[SHELL] Pauses execution for a specified duration in seconds.\n\n" + "Use this to wait for active background operations to progress." +) + + +class WaitTool(ArtemisTool): + """Universal tool for pausing execution for a specified duration in seconds.""" + + def __init__(self, category: ToolCategory = "system"): + super().__init__( + name="wait", + description=WAIT_DOCSTRING, + args_schema=WaitArgs, + category=category, + ) + + # pylint: disable=too-many-arguments,too-many-positional-arguments + async def execute( + self, + driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument + ctx: ArtemisContext | None = None, # pylint: disable=unused-argument + seconds: int | None = None, + **kwargs: Any, + ) -> str: + sec = ( + seconds + if seconds is not None + else (kwargs.get("seconds") or kwargs.get("Seconds") or 1) + ) + if sec < 1: + sec = 1 + elif sec > 60: + sec = 60 + + logger.info(f"Diagnoser waiting for {sec} seconds...") + await asyncio.sleep(sec) + return f"Successfully waited for {sec} seconds." + + +# Universal tool instance & aliases +wait = WaitTool() +Wait = WaitTool +ToolRegistry.register(wait) + + +def get_wait_tool(ctx: ArtemisContext) -> BaseTool: + """Exports wait as a LangChain BaseTool.""" + return trace_langchain_tool(wait.to_langchain_tool(ctx), ctx) + + +wait_wrapper = ToolWrapper( + tool_fn_getter=get_wait_tool, + on_success_fn=lambda seconds: f"Waited {seconds}s", + on_failure_fn=lambda err: f"Wait failed: {err}", +) +export_wait_tool = get_wait_tool diff --git a/artemis/utils/app_launch_utils.py b/artemis/utils/app_launch_utils.py new file mode 100644 index 0000000..94a58a9 --- /dev/null +++ b/artemis/utils/app_launch_utils.py @@ -0,0 +1,255 @@ +# 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. + +"""Utilities for handling app locking and initial app launch logic.""" + +import asyncio + +from artemis.context import AppLaunchResult, ArtemisContext +from artemis.controllers.platform_specific_commands_controller import ( + get_adb_device, + get_current_foreground_package_async, +) +from artemis.controllers.unified_controller import UnifiedMobileController +from artemis.data_engine.trace import TraceSpan +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +SYSTEM_OVERLAYS_VALID_FOR_LAUNCH = { + "com.google.android.permissioncontroller", + "com.android.permissioncontroller", + "com.google.android.packageinstaller", + "com.android.packageinstaller", +} + + +def get_focused_task_package(ctx: ArtemisContext) -> str | None: + """Inspect the active/focused task stack to find the package affinity of the top task. + + Uses adb shell dumpsys activity activities to extract this info. + """ + try: + device = get_adb_device(ctx) + output = str(device.shell("dumpsys activity activities | grep topDisplayFocusedRootTask")) + if "topDisplayFocusedRootTask=" in output: + segment = output.split("topDisplayFocusedRootTask=")[-1] + if "A=" in segment: + affinity = segment.split("A=")[-1].split("}")[0].split()[0].strip() + if ":" in affinity: + return affinity.split(":")[-1] + return affinity + if "{" in segment and "}" in segment: + braces_content = segment.split("{")[-1].split("}")[0] + tokens = braces_content.split() + for token in tokens: + if "." in token and not token.startswith("#") and not token.isdigit(): + if ":" in token: + return token.split(":")[-1] + return token + return None + except Exception as e: + logger.debug(f"Failed to retrieve focused task package via dumpsys: {e}") + return None + + +async def _poll_for_app_ready( + ctx: ArtemisContext, + app_package: str, + max_poll_seconds: int = 15, + poll_interval: float = 1.0, +) -> tuple[bool, str | None]: + """Poll for app to be ready after launch. + + Treats mCurrentFocus=null as a loading state and keeps polling. + Only fails if we get a different (non-null) package or timeout. + + Args: + ctx: Mobile use context + app_package: Expected package name + max_poll_seconds: Maximum time to poll (default: 15s) + poll_interval: Time between polls (default: 1s) + + Returns: + Tuple of (success: bool, error_message: str | None) + """ + polls = int(max_poll_seconds / poll_interval) + + for i in range(polls): + current_package = await get_current_foreground_package_async(ctx) + + if current_package == app_package: + logger.success(f"App {app_package} is ready (took ~{i * poll_interval:.1f}s)") + return True, None + + if current_package in SYSTEM_OVERLAYS_VALID_FOR_LAUNCH: + task_package = get_focused_task_package(ctx) + if task_package == app_package: + logger.success( + f"App {app_package} launch succeeded (system overlay" + f" '{current_package}' focused, but top task stack affinity" + f" matches expected app package '{task_package}', took" + f" ~{i * poll_interval:.1f}s)" + ) + return True, None + + if current_package is None: + logger.debug(f"Poll {i + 1}/{polls}: App loading (mCurrentFocus=null)...") + else: + logger.debug( + f"Poll {i + 1}/{polls}: Wrong app in foreground (expected" + f" '{app_package}', got '{current_package}'). Still waiting..." + ) + + if i < polls - 1: + await asyncio.sleep(poll_interval) + + current_package = await get_current_foreground_package_async(ctx) + error_msg = ( + f"Timeout waiting for {app_package} to load after {max_poll_seconds}s. " + f"Current foreground: {current_package}" + ) + logger.error(error_msg) + return False, error_msg + + +async def launch_app_with_retries( + ctx: ArtemisContext, + app_package: str, + max_retries: int = 3, + max_poll_seconds: int = 15, +) -> tuple[bool, str | None]: + """Launch an app with retry logic and smart polling. + + Args: + ctx: Mobile use context + app_package: Package name (Android) to launch + max_retries: Maximum number of launch attempts (default: 3) + max_poll_seconds: Maximum time to wait for app to load per attempt + (default: 15s) + + Returns: + Tuple of (success: bool, error_message: str | None) + """ + + for attempt in range(1, max_retries + 1): + logger.info(f"Launch attempt {attempt}/{max_retries} for app {app_package}") + + with TraceSpan( + name=f"Launch Attempt {attempt}", + trace_type="span", + ctx=ctx, + ) as span: + span.payload = {"attempt": attempt, "app_package": app_package} + + controller = UnifiedMobileController(ctx) + if attempt > 1: + logger.warning( + f"Attempt {attempt - 1} failed. Force stopping" + f" '{app_package}' to clear frozen state before retrying..." + ) + await controller.terminate_app(app_package) + await asyncio.sleep(1.0) + + launch_success = await controller.launch_app(app_package) + if not launch_success: + error_msg = f"Failed to execute launch command for {app_package}" + logger.error(error_msg) + span.status = "failed" + span.error = error_msg + if attempt == max_retries: + return False, error_msg + await asyncio.sleep(2) + continue + + await asyncio.sleep(1) + + success, error_msg = await _poll_for_app_ready(ctx, app_package, max_poll_seconds) + + if success: + span.status = "success" + span.result = "App is ready" + return True, None + + span.status = "failed" + span.error = error_msg + + if attempt < max_retries: + logger.warning(f"Attempt {attempt} failed: {error_msg}. Retrying...") + await asyncio.sleep(1) + + error_msg = f"Failed to launch {app_package} after {max_retries} attempts" + logger.error(error_msg) + return False, error_msg + + +async def _handle_initial_app_launch( + ctx: ArtemisContext, + locked_app_package: str, +) -> AppLaunchResult: + """Handle initial app launch verification and launching if needed. + + If locked_app_package is set: + 1. Check if the app is already in the foreground + 2. If not, attempt to launch it (with retries) + 3. Return status with success/error information + + Args: + ctx: Mobile use context + locked_app_package: Package name (Android) to lock to + + Returns: + AppLaunchResult with launch status and error information + """ + if not locked_app_package: + error_msg = f"Invalid locked_app_package: '{locked_app_package}'" + logger.error(error_msg) + return AppLaunchResult( + locked_app_package=locked_app_package, + locked_app_initial_launch_success=False, + locked_app_initial_launch_error=error_msg, + ) + + logger.info(f"Starting initial app launch for package: {locked_app_package}") + + try: + current_package = await get_current_foreground_package_async(ctx) + logger.info(f"Current foreground app: {current_package}") + + if current_package == locked_app_package: + logger.info(f"App {locked_app_package} is already in foreground") + return AppLaunchResult( + locked_app_package=locked_app_package, + locked_app_initial_launch_success=True, + locked_app_initial_launch_error=None, + ) + + logger.info(f"App {locked_app_package} not in foreground, attempting to launch") + success, error_msg = await launch_app_with_retries(ctx, locked_app_package) + + return AppLaunchResult( + locked_app_package=locked_app_package, + locked_app_initial_launch_success=success, + locked_app_initial_launch_error=error_msg, + ) + + except Exception as e: + error_msg = f"Exception during initial app launch: {str(e)}" + logger.error(error_msg) + return AppLaunchResult( + locked_app_package=locked_app_package, + locked_app_initial_launch_success=False, + locked_app_initial_launch_error=error_msg, + ) diff --git a/artemis/utils/cli_helpers.py b/artemis/utils/cli_helpers.py new file mode 100644 index 0000000..f7056b2 --- /dev/null +++ b/artemis/utils/cli_helpers.py @@ -0,0 +1,38 @@ +# 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 sys + +from adbutils import AdbClient +from rich.console import Console + + +def display_device_status(console: Console, adb_client: AdbClient | None = None): + """Checks for connected devices and displays the status.""" + console.print("\n[bold]๐Ÿ“ฑ Device Status[/bold]") + devices = None + if adb_client is not None: + devices = adb_client.device_list() + if devices: + console.print("โœ… [bold green]Android device(s) connected:[/bold green]") + for device in devices: + console.print(f" - {device.serial}") + else: + console.print("โŒ [bold red]No Android device found.[/bold red]") + command = "emulator -avd " + if sys.platform not in ["win32", "darwin"]: + command = f"./{command}" + console.print( + f"You can start an emulator using a command like: [bold]'{command}'[/bold]" + ) diff --git a/artemis/utils/cli_selection.py b/artemis/utils/cli_selection.py new file mode 100644 index 0000000..8b288c5 --- /dev/null +++ b/artemis/utils/cli_selection.py @@ -0,0 +1,154 @@ +# 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 sys + +import inquirer +from rich.console import Console +from rich.prompt import Prompt +from rich.text import Text + + +def select_provider_and_model( + console: Console, + available_providers: list[str], + available_models: dict, + default_provider: str, + default_model: str, + provider: str | None = None, + model: str | None = None, +) -> tuple[str, str]: + """Interactive selection of LLM provider and model with arrow-key dropdowns when available. + + Args: + console: Rich console for output + available_providers: List of available provider names + available_models: Dict mapping providers to their available models + default_provider: Default provider to use + default_model: Default model to use + provider: Pre-selected provider (optional) + model: Pre-selected model (optional) + + Returns: + Tuple of (selected_provider, selected_model) + """ + final_provider = provider + final_model = model + + # Interactive provider selection + if not final_provider: + console.print("\n๐Ÿค– [bold cyan]LLM Provider Selection[/bold cyan]") + final_provider = _select_from_list( + console=console, + item_type="provider", + choices=available_providers, + default=default_provider, + message="Select LLM provider", + ) + + # Interactive model selection + if not final_model: + console.print(f"\n๐ŸŽฏ [bold green]Model Selection for {final_provider}[/bold green]") + available_model_list = ( + available_models[final_provider] + if final_provider + else available_models[default_provider] + ) + + default_model_for_provider = ( + default_model if default_model in available_model_list else available_model_list[0] + ) + + final_model = _select_from_list( + console=console, + item_type="model", + choices=available_model_list, + default=default_model_for_provider, + message=f"Select model for {final_provider}", + ) + + return final_provider, final_model + + +def _select_from_list( + console: Console, + item_type: str, + choices: list[str], + default: str, + message: str, +) -> str: + """Select an item from a list using arrow keys when available, fallback to numbered selection. + + Args: + console: Rich console for output + item_type: Type of item being selected (for error messages) + choices: List of choices to select from + default: Default choice + message: Message to display in dropdown + + Returns: + Selected choice + """ + # Try arrow-key dropdown if TTY is available, fallback to numbered selection + if sys.stdin.isatty(): + try: + questions = [ + inquirer.List( + "selection", + message=f"{message} (use arrow keys)", + choices=choices, + default=default, + ), + ] + answers = inquirer.prompt(questions) + return answers["selection"] if answers else default + except Exception: + # Fallback to numbered selection + return _numbered_selection(console, item_type, choices, default) + else: + return _numbered_selection(console, item_type, choices, default) + + +def _numbered_selection(console: Console, item_type: str, choices: list[str], default: str) -> str: + """Fallback numbered selection when arrow keys aren't available.""" + choices_text = "\n".join([f" {i + 1}. {choice}" for i, choice in enumerate(choices)]) + console.print(f"Available {item_type}s:\n{choices_text}") + + default_idx = choices.index(default) + 1 if default in choices else 1 + + while True: + choice = Prompt.ask( + f"Select {item_type} (1-{len(choices)}) or press Enter for default", + default=str(default_idx), + ) + try: + choice_idx = int(choice) - 1 + if 0 <= choice_idx < len(choices): + return choices[choice_idx] + else: + console.print("[red]Invalid choice. Please try again.[/red]") + except ValueError: + console.print("[red]Please enter a number.[/red]") + + +def display_llm_config(console: Console, provider: str, model: str) -> None: + """Display the selected LLM configuration with colors.""" + config_text = Text() + config_text.append("๐Ÿค– LLM Configuration: ", style="bold white") + config_text.append("Provider: ", style="white") + config_text.append(f"{provider}", style="bold cyan") + config_text.append(" | Model: ", style="white") + config_text.append(f"{model}", style="bold green") + + console.print(config_text) diff --git a/artemis/utils/conversations.py b/artemis/utils/conversations.py new file mode 100644 index 0000000..c7a2058 --- /dev/null +++ b/artemis/utils/conversations.py @@ -0,0 +1,45 @@ +# 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. + +from typing import TypeGuard + +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage + + +def is_ai_message(message: BaseMessage) -> TypeGuard[AIMessage]: + return isinstance(message, AIMessage) + + +def is_human_message(message: BaseMessage) -> TypeGuard[HumanMessage]: + return isinstance(message, HumanMessage) + + +def is_tool_message(message: BaseMessage) -> TypeGuard[ToolMessage]: + return isinstance(message, ToolMessage) + + +def is_tool_for_name(tool_message: ToolMessage, name: str) -> bool: + return tool_message.name == name + + +def get_screenshot_message_for_llm(screenshot_base64: str): + prefix = "" if screenshot_base64.startswith("data:image") else "data:image/jpeg;base64," + return HumanMessage( + content=[ + { + "type": "image_url", + "image_url": {"url": f"{prefix}{screenshot_base64}"}, + } + ] + ) diff --git a/artemis/utils/coordinates.py b/artemis/utils/coordinates.py new file mode 100644 index 0000000..eb779a7 --- /dev/null +++ b/artemis/utils/coordinates.py @@ -0,0 +1,344 @@ +# 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 copy +import json +import re +from typing import Any + + +def parse_swipe_parameters( + args: Any, + default_duration: int | None = 400, +) -> tuple[str | None, str | list[int] | None, int | None]: + """Parses swipe/drag/scroll parameters from various possible input formats into a standardized tuple: + + (kind, target, duration) where kind is 'direction' or 'coords' or None. + If 'coords', target is [start_x, start_y, end_x, end_y] in 0-1000 scale. + If 'direction', target is 'up' | 'down' | 'left' | 'right'. + """ + if args is None: + return "direction", "up", default_duration + + # Extract raw dictionary if args is a dict + if isinstance(args, dict): + direction = args.get("direction") + start = args.get("start") + end = args.get("end") + coordinates = args.get("coordinates") or args.get("normalized_coordinates") + action = args.get("action") + gesture = args.get("gesture") + duration = ( + args.get("duration") + or args.get("duration_ms") + or args.get("time_in_ms") + or default_duration + ) + else: + direction = None + start = None + end = None + coordinates = None + action = args + gesture = None + duration = default_duration + + try: + duration_int = int(duration) if duration is not None else default_duration + except (ValueError, TypeError): + duration_int = default_duration + + def _parse_pt(p: Any) -> list[int] | None: + if isinstance(p, (list, tuple)) and len(p) >= 2: + try: + return [int(float(p[0])), int(float(p[1]))] + except Exception: + pass + if isinstance(p, str): + nums = re.findall(r"-?\d+(?:\.\d+)?", p) + if len(nums) >= 2: + return [int(float(nums[0])), int(float(nums[1]))] + return None + + # 1. Check start and end + if start is not None and end is not None: + pt_s = _parse_pt(start) + pt_e = _parse_pt(end) + if pt_s and pt_e: + return "coords", [pt_s[0], pt_s[1], pt_e[0], pt_e[1]], duration_int + + # 2. Check candidate values in priority order + candidates = [coordinates, action, gesture, direction] + for c in candidates: + if c is None: + continue + # If list of 4 numbers + if isinstance(c, (list, tuple)): + if len(c) == 4: + try: + return "coords", [int(float(x)) for x in c], duration_int + except Exception: + pass + elif ( + len(c) == 2 and isinstance(c[0], (list, tuple)) and isinstance(c[1], (list, tuple)) + ): + pt1 = _parse_pt(c[0]) + pt2 = _parse_pt(c[1]) + if pt1 and pt2: + return "coords", [pt1[0], pt1[1], pt2[0], pt2[1]], duration_int + + # If string + if isinstance(c, str): + s = c.strip() + # Try JSON parse first + try: + parsed = json.loads(s) + if isinstance(parsed, (list, tuple)): + if len(parsed) == 4: + return "coords", [int(float(x)) for x in parsed], duration_int + elif len(parsed) == 2: + pt1 = _parse_pt(parsed[0]) + pt2 = _parse_pt(parsed[1]) + if pt1 and pt2: + return "coords", [pt1[0], pt1[1], pt2[0], pt2[1]], duration_int + except Exception: + pass + + # Regex extract numbers + nums = re.findall(r"-?\d+(?:\.\d+)?", s) + if len(nums) == 4: + return "coords", [int(float(x)) for x in nums], duration_int + + # Check direction string (only if no digits) + s_lower = s.lower() + if not any(char.isdigit() for char in s): + for d in ("up", "down", "left", "right"): + if d in s_lower: + return "direction", d, duration_int + + return None, None, duration_int + + +def normalize_point(x: int, y: int, width: int, height: int) -> list[int]: + """Convert physical pixels to normalized 0-1000 coordinates.""" + nx = int(round(x * 1000.0 / width)) + ny = int(round(y * 1000.0 / height)) + # Clamp to [0, 1000] to be safe + return [max(0, min(1000, nx)), max(0, min(1000, ny))] + + +def denormalize_point(nx: int, ny: int, width: int, height: int) -> list[int]: + """Convert normalized 0-1000 coordinates to physical pixels.""" + x = int(round(nx * width / 1000.0)) + y = int(round(ny * height / 1000.0)) + # Clamp to device bounds + return [max(0, min(width - 1, x)), max(0, min(height - 1, y))] + + +def normalize_action_dict(action: dict[str, Any], width: int, height: int) -> dict[str, Any]: + """Recursively search and normalize coordinate fields inside an action dictionary.""" + if not isinstance(action, dict): + return action + + action_copy = copy.deepcopy(action) + + # Normalize 'coordinates' if present and valid + coords = action_copy.get("coordinates") + if coords and isinstance(coords, list): + if len(coords) == 2: + action_copy["coordinates"] = normalize_point(coords[0], coords[1], width, height) + elif len(coords) == 4: + p1 = normalize_point(coords[0], coords[1], width, height) + p2 = normalize_point(coords[2], coords[3], width, height) + action_copy["coordinates"] = p1 + p2 + + # Normalize 'start_coordinates' and 'end_coordinates' for generic formats + for key in ["start_coordinates", "end_coordinates", "target"]: + val = action_copy.get(key) + if val and isinstance(val, list) and len(val) == 2: + action_copy[key] = normalize_point(val[0], val[1], width, height) + + return action_copy + + +def normalize_any_structure(obj: Any, width: int, height: int) -> Any: + """Recursively searches lists and dicts to normalize any action coordinate schemas.""" + if isinstance(obj, list): + return [normalize_any_structure(item, width, height) for item in obj] + elif isinstance(obj, dict): + # If it looks like a single action item + if "action" in obj or "name" in obj: + return normalize_action_dict(obj, width, height) + # Otherwise recurse down keys + return {k: normalize_any_structure(v, width, height) for k, v in obj.items()} + return obj + + +def normalize_step_actions(step_dict: dict[str, Any]) -> dict[str, Any]: + """Normalizes step action representation between action_taken and generic_tools, + + and ensures explicit normalized_coordinates (0-1000 scale) are injected for + UI rendering. + """ + if not isinstance(step_dict, dict): + return step_dict + + extra = step_dict.get("extra_metadata") or {} + width = extra.get("width") or 1080 + height = extra.get("height") or 2400 + + def _enrich_action_item(item: dict[str, Any]) -> dict[str, Any]: + if not isinstance(item, dict): + return item + coords = item.get("coordinates") + act_name = str(item.get("action") or item.get("name") or "").lower() + + if ( + "swipe" in act_name or "scroll" in act_name or "drag" in act_name + ) and "normalized_coordinates" not in item: + kind, target, _ = parse_swipe_parameters(item) + if kind == "coords" and isinstance(target, list) and len(target) == 4: + item["normalized_coordinates"] = target + item["normalized_start_coordinates"] = target[:2] + item["normalized_end_coordinates"] = target[2:] + if not coords: + item["coordinates"] = target + elif kind == "direction" and isinstance(target, str): + if target == "up": + norm = [600, 800, 600, 200] + elif target == "down": + norm = [600, 200, 600, 800] + elif target == "left": + norm = [800, 500, 200, 500] + elif target == "right": + norm = [200, 500, 800, 500] + else: + norm = None + if norm: + item["normalized_coordinates"] = norm + item["normalized_start_coordinates"] = norm[:2] + item["normalized_end_coordinates"] = norm[2:] + + coords = item.get("coordinates") + if coords and isinstance(coords, list) and "normalized_coordinates" not in item: + # If coordinates look physical (e.g., target_bounds present or > 1000), normalize using width/height + if item.get("target_bounds") or any( + isinstance(c, (int, float)) and c > 1000 for c in coords + ): + if len(coords) == 2 and all(isinstance(c, (int, float)) for c in coords): + item["normalized_coordinates"] = [ + int(round(coords[0] * 1000.0 / width)), + int(round(coords[1] * 1000.0 / height)), + ] + elif len(coords) == 4 and all(isinstance(c, (int, float)) for c in coords): + nx1, ny1 = ( + int(round(coords[0] * 1000.0 / width)), + int(round(coords[1] * 1000.0 / height)), + ) + nx2, ny2 = ( + int(round(coords[2] * 1000.0 / width)), + int(round(coords[3] * 1000.0 / height)), + ) + item["normalized_coordinates"] = [nx1, ny1, nx2, ny2] + item["normalized_start_coordinates"] = [nx1, ny1] + item["normalized_end_coordinates"] = [nx2, ny2] + else: + item["normalized_coordinates"] = coords + return item + + action_taken = step_dict.get("action_taken") + if isinstance(action_taken, str): + try: + action_taken = json.loads(action_taken) + step_dict["action_taken"] = action_taken + except Exception: + pass + + if isinstance(action_taken, dict): + step_dict["action_taken"] = _enrich_action_item(action_taken) + elif isinstance(action_taken, list): + step_dict["action_taken"] = [ + _enrich_action_item(x) if isinstance(x, dict) else x for x in action_taken + ] + action_taken = step_dict["action_taken"] + + generic_tools = step_dict.get("generic_tools") + if not isinstance(generic_tools, list) or not generic_tools: + return step_dict + + action_tool_names = { + "click", + "click_sequence", + "long_press", + "input_text", + "swipe", + "press_key", + "object_detection", + "manage_app", + "wait_for_delay", + "wait_for_text", + "report_failure_analysis", + "ask_explorer", + "run_adb_command", + "run_short_adb_command", + } + + has_atomic_action_tool = False + for tool in generic_tools: + if not isinstance(tool, dict) or tool.get("type") == "llm_call": + continue + raw_name = str(tool.get("name", "")).lower().replace("_exec_", "").lstrip("_") + if raw_name in action_tool_names: + has_atomic_action_tool = True + if isinstance(action_taken, (dict, list)): + payload = tool.get("payload") or {} + if isinstance(payload, str): + try: + payload = json.loads(payload) + except Exception: + payload = {} + if not isinstance(payload, dict): + payload = {} + args = payload.get("args") or {} + if isinstance(args, str): + try: + args = json.loads(args) + except Exception: + args = {} + if not isinstance(args, dict): + args = {} + + first_act = ( + action_taken[0] + if isinstance(action_taken, list) + and action_taken + and isinstance(action_taken[0], dict) + else (action_taken if isinstance(action_taken, dict) else {}) + ) + # Enrich missing coordinates/target/input from action_taken + if "target" not in args and "coordinates" not in args: + if first_act.get("coordinates"): + args["coordinates"] = first_act["coordinates"] + if first_act.get("args", {}).get("target"): + args["target"] = first_act["args"]["target"] + if "normalized_coordinates" not in args and first_act.get("normalized_coordinates"): + args["normalized_coordinates"] = first_act["normalized_coordinates"] + if "text" not in args and first_act.get("text"): + args["text"] = first_act["text"] + + args = _enrich_action_item(args) + payload["args"] = args + tool["payload"] = payload + + return step_dict diff --git a/artemis/utils/cv_canvas.py b/artemis/utils/cv_canvas.py new file mode 100644 index 0000000..d509b09 --- /dev/null +++ b/artemis/utils/cv_canvas.py @@ -0,0 +1,189 @@ +# 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 json +from pathlib import Path +import cv2 + + +class ImageCanvas: + def __init__(self, img_id: str, intermediate_dir: str): + self.intermediate_dir = intermediate_dir + + pool_path = Path(intermediate_dir) / "intermediate_transforms.json" + if not pool_path.exists(): + raise ValueError(f"Image pool registry file not found at {pool_path}") + + with open(pool_path) as f: + pool = json.load(f) + + if img_id not in pool: + raise ValueError(f"Image ID '{img_id}' not found in the intermediate image pool.") + + entry = pool[img_id] + img_path = entry["path"] + + self.img = cv2.imread(img_path) + if self.img is None: + raise ValueError(f"Failed to read image at {img_path}") + self.width, self.height = self._get_resolution() + self.transform = entry["transform"].copy() + self.annotations = {} + + def _get_resolution(self) -> tuple[int, int]: + height, width = self.img.shape[:2] + return width, height + + def normalized_to_pixel_coords(self, x: float, y: float) -> tuple[int, int]: + """Convert normalized (0-1000) coordinates to pixel coordinates.""" + px = int(round(x * self.width / 1000.0)) + py = int(round(y * self.height / 1000.0)) + px = max(0, min(self.width - 1, px)) + py = max(0, min(self.height - 1, py)) + return px, py + + def draw_dot(self, x: int, y: int, label_idx: int, radius: int = 10) -> "ImageCanvas": + cx, cy = x, y + # Draw nested white and red circles + cv2.circle(self.img, (cx, cy), 2 * radius, (255, 255, 255), -1) + cv2.circle(self.img, (cx, cy), radius, (0, 0, 255), -1) + + # Draws text "V{label_idx}" using cv2.putText with PIL-matching offset: (x + 2*radius + 4, y - 22) + tx = int(x + 2 * radius + 4) + ty = int(y - 22) + text = f"V{label_idx}" + + cv2.putText( + self.img, + text, + (tx, ty), + cv2.FONT_HERSHEY_SIMPLEX, + 1.5, + (255, 255, 255), + 10, + cv2.LINE_AA, + ) + cv2.putText( + self.img, + text, + (tx, ty), + cv2.FONT_HERSHEY_SIMPLEX, + 1.5, + (0, 0, 255), + 3, + cv2.LINE_AA, + ) + + self.annotations[text] = [cx, cy] + return self + + def save(self, final: bool) -> None: + pool_json_path = Path(self.intermediate_dir) / "intermediate_transforms.json" + if pool_json_path.exists(): + with open(pool_json_path) as f: + pool = json.load(f) + else: + pool = {} + + # Auto-generate next sequential image_id (e.g. "img_1") + img_nums = [] + for k in pool.keys(): + if k.startswith("img_"): + try: + img_nums.append(int(k[4:])) + except ValueError: + pass + next_num = max(img_nums) + 1 if img_nums else 1 + image_id = f"img_{next_num}" + + out_dir = Path(self.intermediate_dir) + out_dir.mkdir(parents=True, exist_ok=True) + out_path = out_dir / f"{image_id}.jpg" + + success = cv2.imwrite(str(out_path), self.img) + if not success: + raise ValueError(f"Failed to save image to {out_path}") + + pool[image_id] = { + "image_id": image_id, + "path": str(out_path.resolve()), + "transform": self.transform.copy(), + "is_output": final, + "annotations": self.annotations.copy(), + } + + with open(pool_json_path, "w") as f: + json.dump(pool, f, indent=2) + + def crop(self, x: int, y: int, w: int, h: int): + if w <= 0 or h <= 0: + raise ValueError("Width and height must be strictly positive.") + + # Calculate bounding boxes + orig_h, orig_w = self.img.shape[:2] + x_end = min(x + w, orig_w) + y_end = min(y + h, orig_h) + x = max(0, x) + y = max(0, y) + + self.img = self.img[y:y_end, x:x_end] + self.width, self.height = self._get_resolution() + + # Adjust offsets by considering the current scale factor + self.transform["offset_x"] += x / self.transform["scale_x"] + self.transform["offset_y"] += y / self.transform["scale_y"] + + # Transform coordinates inside self.annotations + new_annotations = {} + new_w = x_end - x + new_h = y_end - y + for label, coord in self.annotations.items(): + px, py = coord + new_px = px - x + new_py = py - y + if 0 <= new_px < new_w and 0 <= new_py < new_h: + new_annotations[label] = [int(new_px), int(new_py)] + self.annotations = new_annotations + + return self + + def resize_by_factor(self, factor_x: float, factor_y: float = None): + if factor_x <= 0: + raise ValueError("Scale factor must be strictly positive.") + if factor_y is None: + factor_y = factor_x + + height, width = self.img.shape[:2] + new_width = int(width * factor_x) + new_height = int(height * factor_y) + + if new_width == 0 or new_height == 0: + raise ValueError("Resulting dimensions must be greater than 0.") + + self.img = cv2.resize(self.img, (new_width, new_height)) + self.width, self.height = self._get_resolution() + self.transform["scale_x"] *= factor_x + self.transform["scale_y"] *= factor_y + + # Transform coordinates inside self.annotations + new_annotations = {} + for label, coord in self.annotations.items(): + px, py = coord + new_px = int(px * factor_x) + new_py = int(py * factor_y) + if 0 <= new_px < new_width and 0 <= new_py < new_height: + new_annotations[label] = [new_px, new_py] + self.annotations = new_annotations + + return self diff --git a/artemis/utils/decorators.py b/artemis/utils/decorators.py new file mode 100644 index 0000000..68f6a41 --- /dev/null +++ b/artemis/utils/decorators.py @@ -0,0 +1,138 @@ +# 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 asyncio +from collections.abc import Awaitable, Callable +from functools import wraps +from typing import Any, TypeVar, cast, overload + +R = TypeVar("R") + + +def wrap_with_callbacks_sync( + fn: Callable[..., R], + *, + before: Callable[..., None] | None = None, + on_success: Callable[[R], None] | None = None, + on_failure: Callable[[Exception], None] | None = None, + suppress_exceptions: bool = False, +) -> Callable[..., R]: + @wraps(fn) + def wrapper(*args: Any, **kwargs: Any) -> R: + if before: + before() + try: + result = fn(*args, **kwargs) + if on_success: + on_success(result) + return result + except Exception as e: + if on_failure: + on_failure(e) + if suppress_exceptions: + return None # type: ignore + raise + + return wrapper + + +def wrap_with_callbacks_async( + fn: Callable[..., Awaitable[R]], + *, + before: Callable[..., None] | None = None, + on_success: Callable[[R], None] | None = None, + on_failure: Callable[[Exception], None] | None = None, + suppress_exceptions: bool = False, +) -> Callable[..., Awaitable[R]]: + @wraps(fn) + async def wrapper(*args: Any, **kwargs: Any) -> R: + if before: + before() + try: + result = await fn(*args, **kwargs) + if on_success: + on_success(result) + return result + except Exception as e: + if on_failure: + on_failure(e) + if suppress_exceptions: + return None # type: ignore + raise + + return wrapper + + +@overload +def wrap_with_callbacks( + fn: Callable[..., Awaitable[R]], + *, + before: Callable[[], None] | None = ..., + on_success: Callable[[R], None] | None = ..., + on_failure: Callable[[Exception], None] | None = ..., + suppress_exceptions: bool = ..., +) -> Callable[..., Awaitable[R]]: ... + + +@overload +def wrap_with_callbacks( + *, + before: Callable[..., None] | None = ..., + on_success: Callable[[Any], None] | None = ..., + on_failure: Callable[[Exception], None] | None = ..., + suppress_exceptions: bool = ..., +) -> Callable[[Callable[..., R]], Callable[..., R]]: ... + + +@overload +def wrap_with_callbacks( + fn: Callable[..., R], + *, + before: Callable[[], None] | None = ..., + on_success: Callable[[R], None] | None = ..., + on_failure: Callable[[Exception], None] | None = ..., + suppress_exceptions: bool = ..., +) -> Callable[..., R]: ... + + +def wrap_with_callbacks( + fn: Callable[..., Any] | None = None, + *, + before: Callable[[], None] | None = None, + on_success: Callable[[Any], None] | None = None, + on_failure: Callable[[Exception], None] | None = None, + suppress_exceptions: bool = False, +) -> Any: + def decorator(func: Callable[..., Any]) -> Any: + if asyncio.iscoroutinefunction(func): + return wrap_with_callbacks_async( + cast(Callable[..., Awaitable[Any]], func), + before=before, + on_success=on_success, + on_failure=on_failure, + suppress_exceptions=suppress_exceptions, + ) + else: + return wrap_with_callbacks_sync( + cast(Callable[..., Any], func), + before=before, + on_success=on_success, + on_failure=on_failure, + suppress_exceptions=suppress_exceptions, + ) + + if fn is None: + return decorator + else: + return decorator(fn) diff --git a/artemis/utils/file.py b/artemis/utils/file.py new file mode 100644 index 0000000..f313cbd --- /dev/null +++ b/artemis/utils/file.py @@ -0,0 +1,55 @@ +# 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 json +from pathlib import Path +import re +import shutil +from typing import IO +from unittest.mock import MagicMock + + +def strip_json_comments(text: str) -> str: + text = re.sub(r"//.*?$", "", text, flags=re.MULTILINE) + text = re.sub(r"/\*.*?\*/", "", text, flags=re.DOTALL) + return text + + +def load_jsonc(file: IO) -> dict: + return json.loads(strip_json_comments(file.read())) + + +def create_snapshot(src_dir: Path, snapshot_dir: Path): + """Creates a snapshot of a directory by copying it.""" + if ( + isinstance(src_dir, MagicMock) + or isinstance(snapshot_dir, MagicMock) + or "MagicMock" in str(src_dir) + or "MagicMock" in str(snapshot_dir) + ): + return + if not src_dir.exists(): + return + if snapshot_dir.exists(): + shutil.rmtree(snapshot_dir) + shutil.copytree(src_dir, snapshot_dir) + + +def restore_snapshot(snapshot_dir: Path, dst_dir: Path): + """Restores a directory from a snapshot.""" + if not snapshot_dir.exists(): + raise FileNotFoundError(f"Snapshot directory {snapshot_dir} does not exist.") + if dst_dir.exists(): + shutil.rmtree(dst_dir) + shutil.copytree(snapshot_dir, dst_dir) diff --git a/artemis/utils/image_diff.py b/artemis/utils/image_diff.py new file mode 100644 index 0000000..a53e3f2 --- /dev/null +++ b/artemis/utils/image_diff.py @@ -0,0 +1,335 @@ +# 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 asyncio +import base64 +from collections.abc import Awaitable, Callable +import logging +import math +import time + +import cv2 +import numpy as np + +logger = logging.getLogger(__name__) + + +def check_ui_change( + img_before_bytes, + img_after_bytes, + action_item, + full_screen_threshold=0.0001, + roi_threshold=0.05, +): + """Checks if the UI has changed after an action using image diff. + + Args: + img_before_bytes: Image bytes before action. + img_after_bytes: Image bytes after action. + action_item: Dict containing 'action' and optionally 'coordinates'. + full_screen_threshold: Threshold for full screen change (fraction of + changed pixels). + roi_threshold: Threshold for ROI change (fraction of changed pixels in + ROI). + + Returns: + bool: True if UI changed, False otherwise. + """ + try: + # 1. Decode images + nparr_b = np.frombuffer(img_before_bytes, np.uint8) + nparr_a = np.frombuffer(img_after_bytes, np.uint8) + img_b = cv2.imdecode(nparr_b, cv2.IMREAD_GRAYSCALE) + img_a = cv2.imdecode(nparr_a, cv2.IMREAD_GRAYSCALE) + + if img_b is None or img_a is None: + logger.error("Failed to decode images.") + return True # Fallback to assuming change if we can't decode + + h, w = img_b.shape + action_name = action_item.get("action") + coordinates = action_item.get("coordinates") + + # 2. Masking: Ignore status bar (top 10%) and nav bar (bottom 5%) + top_crop = int(h * 0.10) + bottom_crop = int(h * 0.05) + + content_b = img_b[top_crop : h - bottom_crop, :] + content_a = img_a[top_crop : h - bottom_crop, :] + + # 3. Calculate absolute difference and threshold + diff = cv2.absdiff(content_b, content_a) + blurred = cv2.GaussianBlur(diff, (5, 5), 0) + _, thresh = cv2.threshold(blurred, 25, 255, cv2.THRESH_BINARY) + + non_zero_count = cv2.countNonZero(thresh) + total_pixels = content_b.shape[0] * content_b.shape[1] + global_change_ratio = non_zero_count / total_pixels + + logger.info(f"Action: {action_name} | Global UI change ratio: {global_change_ratio:.4f}") + + # 4. Strategy based on action type + if action_name == "swipe": + # Swipe expects large scale movement + return global_change_ratio > full_screen_threshold + + # For tap and others, check ROI if coordinates are available + if ( + action_name + in [ + "tap", + "long_press_on", + "focus_and_input_text", + "focus_and_clear_text", + ] + and coordinates + and len(coordinates) == 2 + ): + cx, cy = coordinates + # Scale ROI size dynamically based on screen diagonal (reference is 1080x2400) + # and enforce a minimum of 15px to prevent collapsing on low-res screens + ref_diagonal = math.sqrt(1080**2 + 2400**2) + current_diagonal = math.sqrt(w**2 + h**2) + scale_factor = current_diagonal / ref_diagonal + roi_size = max(15, int(50 * scale_factor)) + + # Crop ROI from original images to avoid shifting issues with top_crop + y1 = max(0, cy - roi_size) + y2 = min(h, cy + roi_size) + x1 = max(0, cx - roi_size) + x2 = min(w, cx + roi_size) + + roi_b = img_b[y1:y2, x1:x2] + roi_a = img_a[y1:y2, x1:x2] + + roi_diff = cv2.absdiff(roi_b, roi_a) + roi_blurred = cv2.GaussianBlur(roi_diff, (3, 3), 0) + _, roi_thresh = cv2.threshold(roi_blurred, 20, 255, cv2.THRESH_BINARY) + + roi_non_zero = cv2.countNonZero(roi_thresh) + roi_total = roi_b.shape[0] * roi_b.shape[1] + roi_change_ratio = roi_non_zero / roi_total + + logger.info(f"ROI change ratio at [{cx}, {cy}]: {roi_change_ratio:.4f}") + + # If ROI changed significantly, we assume action worked + if roi_change_ratio > roi_threshold: + return True + + # Fallback to global check if no ROI check was triggered or passed + return global_change_ratio > full_screen_threshold + + except Exception as e: + logger.error(f"Error during image diff: {e}") + return True # Fallback to assuming change on error + + +async def wait_for_screen_stability( + take_screenshot_fn: Callable[[], Awaitable[str]], + max_timeout: float = 1.5, + interval: float = 0.2, + stability_threshold: float = 0.001, +) -> str: + """Wait for the screen to stabilize (nearly zero pixel changes between consecutive frames) + + before returning the latest base64 screenshot. + + Args: + take_screenshot_fn: Async callback function that captures a base64 + screenshot. + max_timeout: Maximum time to wait for settling (in seconds) before + forcing exit. + interval: Polling time step between screenshot samples (in seconds). + stability_threshold: Pixel change ratio below which the screen is + considered settled (e.g. 0.001 for < 0.1%). + + Returns: + str: The final settled base64 screenshot. + """ + logger.info("Starting Dynamic Screen Settling...") + + # 1. Capture the initial frame + try: + last_b64 = await take_screenshot_fn() + last_bytes = base64.b64decode(last_b64) + except Exception as e: + logger.error(f"Failed to capture initial frame for settling: {e}") + # Return fallback immediately on failure + return await take_screenshot_fn() + + start_time = time.time() + while time.time() - start_time < max_timeout: + await asyncio.sleep(interval) + + try: + curr_b64 = await take_screenshot_fn() + curr_bytes = base64.b64decode(curr_b64) + except Exception as e: + logger.warning(f"Failed to capture polling frame: {e}") + continue + + # 2. Decode and compare using OpenCV + try: + nparr_b = np.frombuffer(last_bytes, np.uint8) + nparr_a = np.frombuffer(curr_bytes, np.uint8) + img_b = cv2.imdecode(nparr_b, cv2.IMREAD_GRAYSCALE) + img_a = cv2.imdecode(nparr_a, cv2.IMREAD_GRAYSCALE) + + if img_b is None or img_a is None: + last_b64, last_bytes = curr_b64, curr_bytes + continue + + h, w = img_b.shape + + # Masking: crop status bar (top 10%) and navigation bar (bottom 5%) + # to avoid false positives from notification icons, clocks, or blinking elements + top_crop = int(h * 0.10) + bottom_crop = int(h * 0.05) + + content_b = img_b[top_crop : h - bottom_crop, :] + content_a = img_a[top_crop : h - bottom_crop, :] + + # Calculate absolute difference and filter noise + diff = cv2.absdiff(content_b, content_a) + blurred = cv2.GaussianBlur(diff, (5, 5), 0) + _, thresh = cv2.threshold(blurred, 25, 255, cv2.THRESH_BINARY) + + non_zero_count = cv2.countNonZero(thresh) + total_pixels = content_b.shape[0] * content_b.shape[1] + change_ratio = non_zero_count / total_pixels + + elapsed = time.time() - start_time + logger.info( + f"Screen settle poll: diff ratio = {change_ratio:.4f} (target <" + f" {stability_threshold}) | Elapsed: {elapsed:.2f}s" + ) + + # If the change ratio is below stability threshold, screen has stabilized! + if change_ratio < stability_threshold: + logger.info(f"Screen stabilized after {elapsed:.2f}s (diff: {change_ratio:.4f}).") + return curr_b64 + + last_b64, last_bytes = curr_b64, curr_bytes + + except Exception as ex: + logger.error(f"Error during screen stability image diff calculation: {ex}") + last_b64, last_bytes = curr_b64, curr_bytes + continue + + # 3. Timeout safety fallback (e.g., dynamic playing video/animated loaders) + logger.warning( + f"Screen settle timed out after {max_timeout}s without complete" + " stability. Forcing return of latest frame." + ) + return last_b64 + + +async def wait_for_screen_data_stability( + get_screen_data_fn, + max_timeout: float = 1.5, + interval: float = 0.2, + stability_threshold: float = 0.001, +): + """Wait for the screen to stabilize (nearly zero pixel changes between consecutive frames) + + while fetching both screenshot and XML hierarchy directly + (ScreenDataResponse). + + Args: + get_screen_data_fn: Async callback function that captures a + ScreenDataResponse. + max_timeout: Maximum time to wait for settling (in seconds) before + forcing exit. + interval: Polling time step between screen data samples (in seconds). + stability_threshold: Pixel change ratio below which the screen is + considered settled. + + Returns: + ScreenDataResponse: The final settled ScreenDataResponse. + """ + logger.info("Starting Dynamic Screen Data Settling...") + + # 1. Capture the initial frame + try: + last_data = await get_screen_data_fn() + last_bytes = base64.b64decode(last_data.base64) + except Exception as e: + logger.error(f"Failed to capture initial frame for settling: {e}") + # Return fallback immediately on failure + return await get_screen_data_fn() + + start_time = time.time() + while time.time() - start_time < max_timeout: + await asyncio.sleep(interval) + + try: + curr_data = await get_screen_data_fn() + curr_bytes = base64.b64decode(curr_data.base64) + except Exception as e: + logger.warning(f"Failed to capture polling frame: {e}") + continue + + # 2. Decode and compare using OpenCV + try: + nparr_b = np.frombuffer(last_bytes, np.uint8) + nparr_a = np.frombuffer(curr_bytes, np.uint8) + img_b = cv2.imdecode(nparr_b, cv2.IMREAD_GRAYSCALE) + img_a = cv2.imdecode(nparr_a, cv2.IMREAD_GRAYSCALE) + + if img_b is None or img_a is None: + last_data, last_bytes = curr_data, curr_bytes + continue + + h, w = img_b.shape + + # Masking: crop status bar (top 10%) and navigation bar (bottom 5%) + top_crop = int(h * 0.10) + bottom_crop = int(h * 0.05) + + content_b = img_b[top_crop : h - bottom_crop, :] + content_a = img_a[top_crop : h - bottom_crop, :] + + # Calculate absolute difference and filter noise + diff = cv2.absdiff(content_b, content_a) + blurred = cv2.GaussianBlur(diff, (5, 5), 0) + _, thresh = cv2.threshold(blurred, 25, 255, cv2.THRESH_BINARY) + + non_zero_count = cv2.countNonZero(thresh) + total_pixels = content_b.shape[0] * content_b.shape[1] + change_ratio = non_zero_count / total_pixels + + elapsed = time.time() - start_time + logger.info( + f"Screen data settle poll: diff ratio = {change_ratio:.4f}" + f" (target < {stability_threshold}) | Elapsed: {elapsed:.2f}s" + ) + + # If the change ratio is below stability threshold, screen has stabilized! + if change_ratio < stability_threshold: + logger.info(f"Screen stabilized after {elapsed:.2f}s (diff: {change_ratio:.4f}).") + return curr_data + + last_data, last_bytes = curr_data, curr_bytes + + except Exception as ex: + logger.error(f"Error during screen stability image diff calculation: {ex}") + last_data, last_bytes = curr_data, curr_bytes + continue + + # 3. Timeout safety fallback + logger.warning( + f"Screen data settle timed out after {max_timeout}s without complete" + " stability. Forcing return of latest frame." + ) + return last_data diff --git a/artemis/utils/logger.py b/artemis/utils/logger.py new file mode 100644 index 0000000..a2416f1 --- /dev/null +++ b/artemis/utils/logger.py @@ -0,0 +1,212 @@ +# 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. + +from enum import Enum +import logging +import os +from pathlib import Path +import queue +import sys +import threading + +from colorama import Fore, Style, init + +from artemis.data_engine.context_vars import CURRENT_TRACE_ID + +init(autoreset=True) + + +class LogLevel(Enum): + DEBUG = ("DEBUG", Fore.MAGENTA, "๐Ÿ”") + INFO = ("INFO", Fore.WHITE, "โ„น") + SUCCESS = ("SUCCESS", Fore.GREEN, "โœ“") + WARNING = ("WARNING", Fore.YELLOW, "โš ") + ERROR = ("ERROR", Fore.RED, "โŒ") + CRITICAL = ("CRITICAL", Fore.RED + Style.BRIGHT, "๐Ÿ’ฅ") + + +class DataEngineHandler(logging.Handler): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._log_queue = queue.Queue(maxsize=5000) + self._worker_thread = threading.Thread(target=self._drain_queue, daemon=True) + self._worker_thread.start() + + def _drain_queue(self): + while True: + item = self._log_queue.get() + if item is None: + break + engine_mod = sys.modules.get("artemis.data_engine.engine") + if engine_mod and getattr(engine_mod, "_CURRENT_DATA_ENGINE", None): + _CURRENT_DATA_ENGINE = engine_mod._CURRENT_DATA_ENGINE + if _CURRENT_DATA_ENGINE and _CURRENT_DATA_ENGINE.current_session_id: + record_name, message, level_name, formatted, parent_id = item + _CURRENT_DATA_ENGINE.record_trace( + type="log", + name=record_name, + payload={ + "message": message, + "level": level_name, + "formatted": formatted, + }, + status="success", + parent_trace_id=parent_id, + ) + self._log_queue.task_done() + + def emit(self, record): + engine_mod = sys.modules.get("artemis.data_engine.engine") + if not engine_mod or not getattr(engine_mod, "_CURRENT_DATA_ENGINE", None): + return + _CURRENT_DATA_ENGINE = engine_mod._CURRENT_DATA_ENGINE + + if _CURRENT_DATA_ENGINE and _CURRENT_DATA_ENGINE.current_session_id: + log_entry = self.format(record) + log_level = getattr(record, "log_level", None) + if log_level and hasattr(log_level, "value"): + level_name, _color, _symbol = log_level.value + else: + level_name = record.levelname + + parent_id = CURRENT_TRACE_ID.get() + + try: + self._log_queue.put_nowait( + ( + record.name, + record.getMessage(), + level_name, + log_entry, + parent_id, + ) + ) + except queue.Full as e: + raise RuntimeError( + "DataEngine log queue is full - cannot accept more log" + " entries without dropping records" + ) from e + + +class ArtemisLogger: + def __init__( + self, + name: str, + console_level: str = "INFO", + ): + """Initialize the ARTEMIS logger. + + Args: + name: Logger name (usually __name__) + console_level: Minimum level for console output + """ + self.name = name + self.logger = logging.getLogger(name) + self.logger.setLevel(logging.DEBUG) + + self.logger.handlers.clear() + self.logger.propagate = False + + self._setup_console_handler(console_level) + + # Add DataEngineHandler + data_engine_handler = DataEngineHandler() + data_engine_handler.setLevel(logging.DEBUG) + formatter = logging.Formatter(fmt="%(asctime)s | %(name)s | %(levelname)s | %(message)s") + data_engine_handler.setFormatter(formatter) + self.logger.addHandler(data_engine_handler) + + if os.environ.get("ARTEMIS_MCP_SERVER") == "true": + traces_dir = os.environ.get("ARTEMIS_TRACES_PATH", "traces") + log_path = Path(traces_dir) / "mcp_server.log" + log_path.parent.mkdir(parents=True, exist_ok=True) + file_handler = logging.FileHandler(log_path, encoding="utf-8") + file_handler.setLevel(logging.DEBUG) + file_handler.setFormatter(formatter) + self.logger.addHandler(file_handler) + + def _setup_console_handler(self, level: str): + console_handler = logging.StreamHandler(sys.stderr) + console_handler.setLevel(getattr(logging, level.upper())) + + console_formatter = ColoredFormatter() + console_handler.setFormatter(console_formatter) + + self.logger.addHandler(console_handler) + + def debug(self, message: str, **kwargs): + self.logger.debug(message, extra={"log_level": LogLevel.DEBUG}, **kwargs) + + def info(self, message: str, **kwargs): + self.logger.info(message, extra={"log_level": LogLevel.INFO}, **kwargs) + + def success(self, message: str, **kwargs): + self.logger.info(message, extra={"log_level": LogLevel.SUCCESS}, **kwargs) + + def warning(self, message: str, **kwargs): + self.logger.warning(message, extra={"log_level": LogLevel.WARNING}, **kwargs) + + def error(self, message: str, **kwargs): + self.logger.error(message, extra={"log_level": LogLevel.ERROR}, **kwargs) + + def critical(self, message: str, **kwargs): + self.logger.critical(message, extra={"log_level": LogLevel.CRITICAL}, **kwargs) + + def header(self, message: str, **_kwargs): + separator = "=" * 60 + colored_separator = f"{Fore.CYAN}{separator}{Style.RESET_ALL}" + colored_message = f"{Fore.CYAN}{message}{Style.RESET_ALL}" + + print(colored_separator) + print(colored_message) + print(colored_separator) + self.logger.info(f"\n{separator}\n{message}\n{separator}") + + +class ColoredFormatter(logging.Formatter): + def format(self, record): + log_level = getattr(record, "log_level", LogLevel.INFO) + _level_name, color, symbol = log_level.value + + colored_message = f"{color}{symbol} {record.getMessage()}{Style.RESET_ALL}" + + return colored_message + + +_loggers = {} + + +def get_logger( + name: str, + console_level: str = "INFO", +) -> ArtemisLogger: + """Get or create a logger instance. + + Args: + name: Logger name (usually __name__) + log_file: Path to log file (defaults to logs/{name}.log) + console_level: Minimum level for console output + file_level: Minimum level for file output + enable_file_logging: Whether to enable file logging + + Returns: + ArtemisLogger instance + """ + if name not in _loggers: + _loggers[name] = ArtemisLogger( + name=name, + console_level=console_level, + ) + + return _loggers[name] diff --git a/artemis/utils/media.py b/artemis/utils/media.py new file mode 100644 index 0000000..ea32012 --- /dev/null +++ b/artemis/utils/media.py @@ -0,0 +1,207 @@ +# 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 json +import os +from pathlib import Path +import subprocess +import tempfile + +from artemis.utils.logger import get_logger +from PIL import Image + +logger = get_logger(__name__) + +USE_FFMPEG_GIF = os.environ.get("USE_FFMPEG_GIF", "").lower() in ( + "1", + "true", + "yes", +) + + +def quantize_and_save_gif_from_paths( + image_paths: list[Path], + output_path: Path, + colors: int = 128, + duration: int = 100, +) -> None: + """Create an optimized GIF from image file paths. + + By default uses PIL (loads all frames into memory). + Set USE_FFMPEG_GIF=1 env var to use ffmpeg for memory-efficient streaming. + + Args: + image_paths: List of paths to image files (must be sorted in desired + order) + output_path: Path where the GIF will be saved + colors: Number of colors to use in quantization (lower = smaller file) + duration: Duration of each frame in milliseconds + + Raises: + ValueError: If image_paths list is empty + RuntimeError: If ffmpeg fails (when USE_FFMPEG_GIF is enabled) + """ + if not image_paths: + raise ValueError("image_paths must not be empty") + + if USE_FFMPEG_GIF: + _save_gif_ffmpeg(image_paths, output_path, colors, duration) + else: + _save_gif_pillow(image_paths, output_path, colors, duration) + + +def _save_gif_pillow( + image_paths: list[Path], + output_path: Path, + colors: int, + duration: int, +) -> None: + """Create GIF using PIL (loads all frames into memory).""" + + def frame_generator(): + for path in image_paths: + with Image.open(path) as img: + if img.mode != "RGB": + img = img.convert("RGB") + yield img.quantize(colors=colors, method=2) + + frames = frame_generator() + first_frame = next(frames) + first_frame.save( + output_path, + save_all=True, + append_images=frames, + loop=0, + optimize=True, + duration=duration, + ) + + +def _save_gif_ffmpeg( + image_paths: list[Path], + output_path: Path, + colors: int, + duration: int, +) -> None: + """Create GIF using ffmpeg (memory-efficient streaming).""" + fps = 1000 / duration + + with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: + for i, path in enumerate(image_paths): + f.write(f"file '{path.absolute()}'\n") + # Last file needs no duration (ffmpeg concat demuxer uses it as the final frame) + if i < len(image_paths) - 1: + f.write(f"duration {duration / 1000}\n") + # Repeat last file to ensure it's included (ffmpeg concat demuxer quirk) + f.write(f"file '{image_paths[-1].absolute()}'\n") + concat_file = Path(f.name) + + palette_path = output_path.with_suffix(".palette.png") + + try: + result = subprocess.run( + [ + "ffmpeg", + "-y", + "-f", + "concat", + "-safe", + "0", + "-i", + str(concat_file), + "-vf", + f"palettegen=max_colors={min(colors, 256)}:stats_mode=diff", + str(palette_path), + ], + capture_output=True, + text=True, + ) + if result.returncode != 0: + raise RuntimeError(f"ffmpeg palette generation failed: {result.stderr}") + + result = subprocess.run( + [ + "ffmpeg", + "-y", + "-f", + "concat", + "-safe", + "0", + "-i", + str(concat_file), + "-i", + str(palette_path), + "-lavfi", + f"fps={fps},paletteuse=dither=bayer:bayer_scale=5", + "-loop", + "0", + str(output_path), + ], + capture_output=True, + text=True, + ) + if result.returncode != 0: + raise RuntimeError(f"ffmpeg GIF creation failed: {result.stderr}") + + finally: + concat_file.unlink(missing_ok=True) + if palette_path.exists(): + palette_path.unlink() + + +def create_gif_from_trace_folder(trace_folder_path: Path): + image_files: list[Path] = [] + + for file in trace_folder_path.iterdir(): + if file.suffix == ".jpeg": + image_files.append(file) + + image_files.sort(key=lambda f: int(f.stem)) + + logger.info(f"Found {len(image_files)} images to compile") + + if not image_files: + return + + gif_path = trace_folder_path / "trace.gif" + quantize_and_save_gif_from_paths(image_files, gif_path) + logger.info(f"GIF created at {gif_path}") + + +def remove_images_from_trace_folder(trace_folder_path: Path): + for file in trace_folder_path.iterdir(): + if file.suffix == ".jpeg": + file.unlink() + + +def create_steps_json_from_trace_folder(trace_folder_path: Path): + steps = [] + for file in trace_folder_path.iterdir(): + if file.suffix == ".json": + with open(file, encoding="utf-8", errors="ignore") as f: + json_content = f.read() + steps.append({"timestamp": int(file.stem), "data": json_content}) + + steps.sort(key=lambda f: f["timestamp"]) + + logger.info(f"Found {len(steps)} steps to compile") + + with open(trace_folder_path / "steps.json", "w", encoding="utf-8", errors="ignore") as f: + f.write(json.dumps(steps)) + + +def remove_steps_json_from_trace_folder(trace_folder_path: Path): + for file in trace_folder_path.iterdir(): + if file.suffix == ".json" and file.name != "steps.json": + file.unlink() diff --git a/artemis/utils/notes.py b/artemis/utils/notes.py new file mode 100644 index 0000000..c32a633 --- /dev/null +++ b/artemis/utils/notes.py @@ -0,0 +1,441 @@ +# 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 difflib +import hashlib +import json +import os +from pathlib import Path +import re + +from artemis.data_engine.engine import _CURRENT_DATA_ENGINE + +# Centralized Description and Documentation Constants (Single Source of Truth) +READ_NOTE_DOCSTRING = ( + "[NOTE] Reads a previously saved note from persistent memory by its" + " key.\nUse this as persistent cross-turn memory to retrieve key data, plan" + " cross-turn subgoals, and coordinate step progress.\nYou can optionally" + " specify a line range to read by providing start_line and end_line" + " (1-indexed, inclusive).\nTIP: For large notes, prefer reading a small" + " range of lines at a time (e.g., up to 200 lines). For chronological" + " files, consider reading from back to front to see the latest entries" + " first.\n\nUse the key 'task_plan' to read the active task plan" + " (checklist)." +) +READ_NOTE_ARG_KEY_DESC = "The key of the note to read." + +LIST_NOTES_DOCSTRING = ( + "[NOTE] Lists all note keys currently stored in persistent memory, along" + " with the number of lines in each note.\nUse this to coordinate step" + " progress and recall persistent cross-turn memories." +) + +SAVE_NOTE_DOCSTRING = ( + "[NOTE] Saves a text note to persistent memory with the given key. Use this" + " as persistent cross-turn memory to save key data, plan cross-turn" + " subgoals, and coordinate step progress. If the key already exists, it" + " will be overwritten.\n\nUse the key 'task_plan' to overwrite the original" + " task plan (checklist)." +) +SAVE_NOTE_ARG_KEY_DESC = "The unique key under which to save or overwrite the note." +SAVE_NOTE_ARG_CONTENT_DESC = "The text content to write to the note." + +APPEND_NOTE_DOCSTRING = ( + "[NOTE] Appends text content to an existing note in persistent memory with" + " the given key. Use this as persistent cross-turn memory to record" + " chronological logs, save key data, or coordinate step progress. If the" + " note does not exist, it will be created.\n\nUse the key 'task_plan' to" + " append to the task plan (checklist)." +) +APPEND_NOTE_ARG_KEY_DESC = "The key of the note to append to." +APPEND_NOTE_ARG_CONTENT_DESC = "The text content to append to the note." + +UPDATE_NOTE_DOCSTRING = ( + "[NOTE] Updates a note by replacing a specific target string with" + " replacement content. Use this as persistent cross-turn memory to modify" + " specific planning nodes or step details without overwriting the whole" + " file.\n\nUse the key 'task_plan' to update specific sections of the task" + " plan (checklist)." +) +UPDATE_NOTE_ARG_KEY_DESC = "The key of the note to update." +UPDATE_NOTE_ARG_TARGET_DESC = ( + "The exact target string currently in the note that you wish to replace." +) +UPDATE_NOTE_ARG_REPLACEMENT_DESC = ( + "The new replacement content to insert instead of the target content." +) + + +# Core Business Logic Filesystem Functions + + +def get_notes_dir(base_dir: str | Path) -> Path: + """Gets the absolute path to the notes directory inside the session's base directory.""" + return Path(base_dir) / "notes" + + +def get_note_file_path(base_dir: str | Path, key: str) -> Path: + """Gets the path of a note file by its key, mapping 'task_plan' to 'task_plan.md'. + + Case-insensitively strips any trailing '.md' from the key to prevent + double-extension bugs. + """ + notes_dir = get_notes_dir(base_dir) + # Strip any trailing .md (e.g. 'findings.md' -> 'findings') + clean_key = key[:-3] if key.lower().endswith(".md") else key + if clean_key == "task_plan": + return notes_dir / "task_plan.md" + return notes_dir / f"{clean_key}.md" + + +def list_notes_keys(base_dir: str | Path) -> list[str]: + """Lists all note keys currently stored in persistent memory.""" + notes_dir = get_notes_dir(base_dir) + if not notes_dir.exists(): + return [] + return [f.stem for f in notes_dir.glob("*.md")] + + +def list_notes_info(base_dir: str | Path) -> dict[str, int]: + """Lists all note keys currently stored in persistent memory, along with their line counts.""" + notes_dir = get_notes_dir(base_dir) + if not notes_dir.exists(): + return {} + info = {} + for f in notes_dir.glob("*.md"): + try: + content = f.read_text(encoding="utf-8") + info[f.stem] = len(content.splitlines()) + except Exception: + info[f.stem] = -1 + return info + + +def read_note_content( + base_dir: str | Path, + key: str, + start_line: int | None = None, + end_line: int | None = None, +) -> str: + """Reads the raw content of a note, optionally restricted to a line range (1-indexed, inclusive).""" + file_path = get_note_file_path(base_dir, key) + if not file_path.exists(): + raise FileNotFoundError(f"Note '{key}' not found.") + content = file_path.read_text(encoding="utf-8") + + if start_line is not None or end_line is not None: + lines = content.splitlines() + start = (start_line - 1) if start_line else 0 + end = end_line if end_line else len(lines) + return "\n".join(lines[start:end]) + + return content + + +def save_note_content(base_dir: str | Path, key: str, content: str) -> None: + """Saves or overwrites the raw content of a note.""" + notes_dir = get_notes_dir(base_dir) + notes_dir.mkdir(parents=True, exist_ok=True) + file_path = get_note_file_path(base_dir, key) + + content_before = "" + if key == "task_plan" and file_path.exists(): + try: + content_before = file_path.read_text(encoding="utf-8") + except Exception: + pass + + file_path.write_text(content, encoding="utf-8") + + if key == "task_plan": + try: + record_subgoal_hash_chain(base_dir, content_before, content) + except Exception: + pass + + try: + if _CURRENT_DATA_ENGINE: + _CURRENT_DATA_ENGINE._publish("note_saved", {"key": key, "content": content}) + except Exception: + pass + + +def append_note_content(base_dir: str | Path, key: str, content: str) -> None: + """Appends raw content to a note, inserting a newline if the file is not empty and lacks a trailing newline.""" + notes_dir = get_notes_dir(base_dir) + notes_dir.mkdir(parents=True, exist_ok=True) + file_path = get_note_file_path(base_dir, key) + + prefix = "" + if file_path.exists() and file_path.stat().st_size > 0: + # Open in binary mode for safe and portable backward seek + with open(file_path, "rb") as f: + try: + f.seek(-1, os.SEEK_END) + last_char_bytes = f.read(1) + if last_char_bytes != b"\n": + prefix = "\n" + except Exception: + pass + + with open(file_path, "a", encoding="utf-8") as f: + f.write(prefix + content + "\n") + + +def clean_line(line: str) -> str: + """Cleans a line by lowering case, stripping quotes, removing markdown bullet/checkbox markers, and all whitespace.""" + s = line.lower() + s = s.replace('"', "").replace("'", "") + # Strip markdown bullet patterns: optional start spaces, bullet character (-*+), optional spaces, optional checkbox [ xX/] + s = re.sub(r"^\s*[-*+]\s*(\[[ xX/]\])?\s*", "", s) + return "".join(s.split()) + + +def find_relaxed_match(content: str, target: str) -> tuple[int, int] | None: + """Finds a unique relaxed match (ignoring whitespaces, casing, quotes, and list bullets). + + Returns (start_char_idx, end_char_idx) if a unique match is found, otherwise + None. + """ + note_lines = content.splitlines(keepends=True) + target_lines = target.splitlines() + + clean_note = [clean_line(line) for line in note_lines] + clean_target = [clean_line(line) for line in target_lines if clean_line(line)] + + n = len(clean_note) + m = len(clean_target) + if m == 0 or n < m: + return None + + matches = [] + for i in range(n - m + 1): + window = clean_note[i : i + m] + if window == clean_target: + matches.append((i, i + m)) + + if len(matches) == 1: + start_line, end_line = matches[0] + start_idx = sum(len(line) for line in note_lines[:start_line]) + end_idx = sum(len(line) for line in note_lines[:end_line]) + return start_idx, end_idx + + return None + + +def find_fuzzy_match(content: str, target: str, threshold: float = 0.90) -> tuple[int, int] | None: + """Uses sliding window difflib similarity to find a unique close match above the threshold. + + Returns (start_char_idx, end_char_idx) if a unique close match is found, + otherwise None. + """ + note_lines = content.splitlines(keepends=True) + target_lines = target.splitlines() + + n = len(note_lines) + m = len(target_lines) + if m == 0 or n < m: + return None + + best_match = None + best_ratio = 0.0 + ambiguous = False + + target_clean = "\n".join(clean_line(line) for line in target_lines) + + for i in range(n - m + 1): + window_clean = "\n".join(clean_line(line) for line in note_lines[i : i + m]) + matcher = difflib.SequenceMatcher(None, window_clean, target_clean) + ratio = matcher.ratio() + + if ratio >= threshold: + if ratio > best_ratio + 0.05: + best_ratio = ratio + best_match = (i, i + m) + ambiguous = False + elif abs(ratio - best_ratio) <= 0.05: + ambiguous = True + + if best_match and not ambiguous: + start_line, end_line = best_match + start_idx = sum(len(line) for line in note_lines[:start_line]) + end_idx = sum(len(line) for line in note_lines[:end_line]) + return start_idx, end_idx + + return None + + +def apply_replacement(content: str, start_idx: int, end_idx: int, replacement: str) -> str: + """Applies a replacement, preserving the trailing newline of the original slice if present.""" + original_slice = content[start_idx:end_idx] + actual_replacement = replacement + if original_slice.endswith("\n") and not actual_replacement.endswith("\n"): + actual_replacement += "\n" + return content[:start_idx] + actual_replacement + content[end_idx:] + + +def update_note_content( + base_dir: str | Path, key: str, target: str, replacement: str +) -> str | None: + """Replaces target content in a note using hierarchical exact, relaxed, or fuzzy match. + + Returns a warning message if relaxed/fuzzy match was used, otherwise None. + Raises ValueError if target is not found or if the match is ambiguous/not + unique. + """ + file_path = get_note_file_path(base_dir, key) + if not file_path.exists(): + raise FileNotFoundError(f"Note '{key}' not found.") + + content = file_path.read_text(encoding="utf-8") + new_content_to_write = None + warning = None + + # Phase 1: Exact Match + if target in content: + count = content.count(target) + if count == 1: + new_content_to_write = content.replace(target, replacement) + warning = None + elif count > 1: + raise ValueError( + f"Target string '{target}' is not unique in note '{key}' (found" + f" {count} occurrences). To prevent accidental changes, please" + " provide a longer and more complete target string." + ) + + # Phase 2: Relaxed Match (Whitespace, casing, quotes, checkboxes ignored) + if new_content_to_write is None: + relaxed_range = find_relaxed_match(content, target) + if relaxed_range: + start_idx, end_idx = relaxed_range + new_content_to_write = apply_replacement(content, start_idx, end_idx, replacement) + warning = ( + "We did our best to apply changes despite some inaccuracies." + " Double check if the edit applied is what you intended." + ) + + # Phase 3: Fuzzy Match (difflib similarity) + if new_content_to_write is None: + fuzzy_range = find_fuzzy_match(content, target, threshold=0.90) + if fuzzy_range: + start_idx, end_idx = fuzzy_range + new_content_to_write = apply_replacement(content, start_idx, end_idx, replacement) + warning = ( + "We did our best to apply changes despite some inaccuracies." + " Double check if the edit applied is what you intended." + ) + + # Phase 4: Write and Return, or Not Found Error + if new_content_to_write is not None: + file_path.write_text(new_content_to_write, encoding="utf-8") + if key == "task_plan": + try: + record_subgoal_hash_chain(base_dir, content, new_content_to_write) + except Exception: + pass + return warning + + raise ValueError( + f"Target string '{target}' not found in note '{key}'. Please re-read" + " the note and ensure spelling, spacing, and indentation are close." + ) + + +# Shared Formatting Functions for Tool Outputs + + +def format_read_note_success( + key: str, + content: str, + start_line: int | None = None, + end_line: int | None = None, +) -> str: + """Formats the success message for reading a note.""" + range_str = "" + if start_line is not None or end_line is not None: + range_str = f" (lines {start_line or 1} to {end_line or 'end'})" + return f"Successfully read note '{key}'{range_str}. '{key}' note content:\n{content}" + + +def format_read_note_failure(key: str, error: str) -> str: + """Formats the failure message for reading a note.""" + if "not found" in error.lower(): + return f"Note '{key}' not found in scratchpad." + return f"Failed to read note '{key}': {error}" + + +def format_list_notes_success(notes_info: dict[str, int]) -> str: + """Formats the success message for listing notes with line counts.""" + if not notes_info: + return "No notes saved yet." + lines = [f"- {key} ({count} lines)" for key, count in notes_info.items()] + return "Here are all the notes:\n" + "\n".join(lines) + + +def format_list_notes_failure(error: str) -> str: + """Formats the failure message for listing notes.""" + return f"Failed to list notes: {error}" + + +def record_subgoal_hash_chain( + base_dir: str | Path, content_before: str, content_after: str +) -> None: + """Detects if the active subgoal was renamed/changed, and records it in the subgoal_hash_chain.json.""" + if not content_before or not content_after: + return + + # Helper to get the active subgoal hash from task_plan string + def get_active_hash(task_plan: str) -> str | None: + try: + lines = task_plan.splitlines() + for line in reversed(lines): + if line.strip().startswith("- [/]"): + text = line.strip()[5:].strip() + return hashlib.md5(text.encode("utf-8")).hexdigest() + + # Fallback to first pending [ ] if no [/] + for line in lines: + if line.strip().startswith("- [ ]"): + text = line.strip()[5:].strip() + return hashlib.md5(text.encode("utf-8")).hexdigest() + except Exception: + pass + return None + + old_active = get_active_hash(content_before) + new_active = get_active_hash(content_after) + + if old_active and new_active and old_active != new_active: + notes_dir = Path(base_dir) / "notes" + notes_dir.mkdir(parents=True, exist_ok=True) + chain_path = notes_dir / "subgoal_hash_chain.json" + + chain = {} + if chain_path.exists(): + try: + chain = json.loads(chain_path.read_text(encoding="utf-8")) + except Exception: + pass + + # Record mapping: old -> new + chain[old_active] = new_active + + try: + chain_path.write_text( + json.dumps(chain, indent=2, ensure_ascii=False), + encoding="utf-8", + ) + except Exception: + pass diff --git a/artemis/utils/ocr_api.py b/artemis/utils/ocr_api.py new file mode 100644 index 0000000..a338aa1 --- /dev/null +++ b/artemis/utils/ocr_api.py @@ -0,0 +1,90 @@ +# 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 os +from typing import Any + +import httpx +from artemis.config import settings +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) +_HTTP_CLIENT: httpx.AsyncClient | None = None + + +def get_http_client() -> httpx.AsyncClient: + global _HTTP_CLIENT + if _HTTP_CLIENT is None or _HTTP_CLIENT.is_closed: + _HTTP_CLIENT = httpx.AsyncClient(timeout=30.0) + return _HTTP_CLIENT + + +async def perform_ocr( + screenshot_b64: str, client: httpx.AsyncClient = None +) -> list[dict[str, Any]]: + """Calls Google Vision API for OCR using a persistent client to reuse TLS handshakes. + + Args: + screenshot_b64: Base64 encoded screenshot image. + + Returns: + A list of dictionaries containing detected text and position vertices. + """ + ocr_secret = settings.get_api_key("ocr") + api_key = ( + (ocr_secret.get_secret_value() if ocr_secret else None) + or os.environ.get("OCR_API_KEY") + or os.environ.get("VISION_API_KEY") + or os.environ.get("API_KEY") + or os.environ.get("GOOGLE_API_KEY") + or "API_KEY" + ) + + url = f"https://vision.googleapis.com/v1/images:annotate?key={api_key}" + headers = {"Content-Type": "application/json"} + data = { + "requests": [ + { + "image": {"content": screenshot_b64}, + "features": [{"type": "TEXT_DETECTION"}], + } + ] + } + + try: + active_client = client if client is not None else get_http_client() + response = await active_client.post(url, json=data, headers=headers, timeout=30.0) + return _parse_ocr_response(response) + except Exception as e: + logger.error(f"Failed to call Vision API for OCR: {e}") + raise e + + +def _parse_ocr_response(response: httpx.Response) -> list[dict[str, Any]]: + if response.status_code == 200: + res_json = response.json() + responses = res_json.get("responses", []) + if responses and "textAnnotations" in responses[0]: + annotations = responses[0]["textAnnotations"] + results = [] + # Skip the first annotation (index 0) as it is the full-screen combined text + for ann in annotations[1:]: + desc = ann.get("description", "") + vertices = ann.get("boundingPoly", {}).get("vertices", []) + results.append({"text": desc, "position": vertices}) + return results + else: + return [] + else: + raise Exception(f"Vision API returned status {response.status_code}: {response.text}") diff --git a/artemis/utils/ocr_xml_fusion.py b/artemis/utils/ocr_xml_fusion.py new file mode 100644 index 0000000..57d69db --- /dev/null +++ b/artemis/utils/ocr_xml_fusion.py @@ -0,0 +1,318 @@ +# 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 difflib +from io import BytesIO +from typing import Any + +from artemis.utils.logger import get_logger +from artemis.utils.ui_filter import _parse_bounds +from PIL import Image + +logger = get_logger(__name__) + +PLACEHOLDERS = {"search", "enter text", "placeholder"} + + +def _is_low_value_text(text: str) -> bool: + """Check if the text is low value or a placeholder for virtual nodes.""" + if not text: + return True + + text_strip = text.strip() + + # 1. Extremely short character check + if len(text_strip) == 1: + # Check if it's a Chinese character + is_chinese = "\u4e00" <= text_strip <= "\u9fff" + # Check if it's a digit + is_digit = text_strip.isdigit() + + # If it's neither Chinese nor digit, consider it low value (likely noise) + if not (is_chinese or is_digit): + return True + + # 2. Common placeholders + if text_strip.lower() in PLACEHOLDERS: + return True + + # 3. Default markers from tools + if text_strip in {"[Icon]", "[Image]", "[Picture]"}: + return True + + return False + + +def _parse_ocr_position(position: list[dict[str, int]]) -> dict[str, int]: + """Convert OCR vertices to [left, top, right, bottom].""" + if not position: + return {"left": 0, "top": 0, "right": 0, "bottom": 0} + left = min(v.get("x", 0) for v in position) + top = min(v.get("y", 0) for v in position) + right = max(v.get("x", 0) for v in position) + bottom = max(v.get("y", 0) for v in position) + return {"left": left, "top": top, "right": right, "bottom": bottom} + + +def _calculate_overlap_ratio(ocr_bounds: dict[str, int], xml_bounds: dict[str, int]) -> float: + """Calculate how much of the OCR box is inside the XML node (Intersection / OCR_Area).""" + inter_left = max(ocr_bounds["left"], xml_bounds["left"]) + inter_top = max(ocr_bounds["top"], xml_bounds["top"]) + inter_right = min(ocr_bounds["right"], xml_bounds["right"]) + inter_bottom = min(ocr_bounds["bottom"], xml_bounds["bottom"]) + + if inter_left >= inter_right or inter_top >= inter_bottom: + return 0.0 + + inter_area = (inter_right - inter_left) * (inter_bottom - inter_top) + ocr_area = (ocr_bounds["right"] - ocr_bounds["left"]) * ( + ocr_bounds["bottom"] - ocr_bounds["top"] + ) + + if ocr_area == 0: + return 0.0 + + return inter_area / ocr_area + + +def _detect_status_bar_height(xml_hierarchy: list[dict[str, Any]], screen_height: int) -> int: + """Try to detect status bar height from XML, fallback to 4% of screen height.""" + for node in xml_hierarchy: + package = node.get("package") or node.get("packageName") + bounds_str = node.get("bounds") + + # Look for systemui at the top + if package == "com.android.systemui": + bounds = _parse_bounds(bounds_str) + if bounds and bounds["top"] == 0: + return bounds["bottom"] + + return int(screen_height * 0.04) + + +def _crop_image_remove_status_bar(screenshot_b64: str, crop_height: int) -> tuple[str, int, int]: + """Crop the top part of the image. + + Returns cropped b64, original width, original height. + """ + img_data = base64.b64decode(screenshot_b64) + img = Image.open(BytesIO(img_data)) + width, height = img.size + + if crop_height >= height: + logger.warning("Crop height is greater than image height. Not cropping.") + return screenshot_b64, width, height + + # Crop: left, top, right, bottom + cropped_img = img.crop((0, crop_height, width, height)) + + buffered = BytesIO() + cropped_img.save(buffered, format="PNG") + cropped_b64 = base64.b64encode(buffered.getvalue()).decode("utf-8") + + return cropped_b64, width, height + + +def _map_coordinates_back( + ocr_results: list[dict[str, Any]], status_bar_height: int +) -> list[dict[str, Any]]: + """Map OCR coordinates back to original screen coordinates.""" + mapped_results = [] + for res in ocr_results: + text = res.get("text", "") + position = res.get("position", []) + + mapped_position = [] + for v in position: + mapped_position.append({"x": v.get("x", 0), "y": v.get("y", 0) + status_bar_height}) + + mapped_results.append({"text": text, "position": mapped_position}) + return mapped_results + + +def fuse_ocr_with_xml( + xml_hierarchy: list[dict[str, Any]], ocr_results: list[dict[str, Any]] +) -> list[dict[str, Any]]: + """Fuse OCR results with XML hierarchy with advanced line-grouping aggregation. + + Assumes ocr_results coordinates are already mapped back to full screen. + """ + # 0. Detect screen dimensions to calculate screen area + max_width = 1080 + max_height = 2400 + for node in xml_hierarchy: + bounds_str = node.get("bounds") + xml_bounds = _parse_bounds(bounds_str) + if xml_bounds: + max_width = max(max_width, xml_bounds["right"]) + max_height = max(max_height, xml_bounds["bottom"]) + screen_area = max_width * max_height + + # 1. Parse OCR positions and keep in list + ocr_with_bounds = [] + for ocr in ocr_results: + bounds = _parse_ocr_position(ocr.get("position", [])) + ocr_with_bounds.append({"ocr": ocr, "bounds": bounds}) + + # 2. Match OCR to XML nodes (Banded Overlap + Giant Node Exclusion) + xml_to_ocr_map = {} + unmatched_ocrs = [] + + for idx, ocr_item in enumerate(ocr_with_bounds): + ocr_bounds = ocr_item["bounds"] + ocr_text = ocr_item["ocr"].get("text", "").strip() + if not ocr_text: + continue + + band1_candidates = [] # list of (node_idx, area) + band2_candidates = [] # list of (node_idx, area) + + for node_idx, node in enumerate(xml_hierarchy): + bounds_str = node.get("bounds") + xml_bounds = _parse_bounds(bounds_str) + if not xml_bounds: + continue + + node_area = (xml_bounds["right"] - xml_bounds["left"]) * ( + xml_bounds["bottom"] - xml_bounds["top"] + ) + node_class = str(node.get("className") or node.get("class") or "") + + # Exclude Giant Nodes (unless text-centric class) + is_text_class = any( + tc in node_class.lower() for tc in ["text", "button", "edit", "search", "input"] + ) + if node_area > screen_area * 0.1 and not is_text_class: + continue # Exclude giant container + + overlap = _calculate_overlap_ratio(ocr_bounds, xml_bounds) + + if overlap >= 0.9: + band1_candidates.append((node_idx, node_area)) + elif overlap >= 0.7: + band2_candidates.append((node_idx, node_area)) + + best_node_idx = None + if band1_candidates: + # Sort by area ascending (prefer smaller/more specific nodes) + band1_candidates.sort(key=lambda x: x[1]) + best_node_idx = band1_candidates[0][0] + elif band2_candidates: + band2_candidates.sort(key=lambda x: x[1]) + best_node_idx = band2_candidates[0][0] + + if best_node_idx is not None: + if best_node_idx not in xml_to_ocr_map: + xml_to_ocr_map[best_node_idx] = [] + xml_to_ocr_map[best_node_idx].append(ocr_item) + else: + unmatched_ocrs.append(ocr_item) + + # 3. Advanced Aggregation for Matched Nodes (Line Grouping) + for node_idx, matched_list in xml_to_ocr_map.items(): + node = xml_hierarchy[node_idx] + xml_text = str(node.get("text") or "").strip() + + # Group items into lines to handle vertical misalignment + lines = [] + matched_list.sort(key=lambda x: x["bounds"]["top"]) + + for item in matched_list: + b = item["bounds"] + placed = False + for line in lines: + # If vertical center is close to line's average vertical center + line_center = sum(x["bounds"]["top"] + x["bounds"]["bottom"] for x in line) / ( + 2 * len(line) + ) + item_center = (b["top"] + b["bottom"]) / 2 + if abs(line_center - item_center) < 10: # Stricter 10 pixel tolerance for same line + line.append(item) + placed = True + break + if not placed: + lines.append([item]) + + # Sort lines from top to bottom, and items within line from left to right + lines.sort(key=lambda row: sum(x["bounds"]["top"] for x in row) / len(row)) + for line in lines: + line.sort(key=lambda x: x["bounds"]["left"]) + + # Flatten and join text + sorted_items = [item for line in lines for item in line] + aggregated_ocr_text = " ".join(item["ocr"].get("text", "").strip() for item in sorted_items) + + # Discard Strategy (Deduplication) + # We synthesize the full text to check for redundancy against XML text + discard = False + if xml_text and (aggregated_ocr_text in xml_text): + discard = True + if xml_text and not discard: + similarity = difflib.SequenceMatcher( + None, xml_text.lower(), aggregated_ocr_text.lower() + ).ratio() + if similarity > 0.8: + discard = True + + if not discard: + # Instead of a single flat string, we list OCR elements line by line + ocr_elements = [] + for line in lines: + current_text = "" + current_bounds = None + + for item in line: + txt = item["ocr"].get("text", "").strip() + b = item["bounds"] + + if not current_bounds: + current_bounds = b.copy() + current_text = txt + else: + # Strict horizontal tolerance: if within 30 pixels, merge them + if b["left"] - current_bounds["right"] < 30: + current_bounds["right"] = b["right"] + current_bounds["bottom"] = max(current_bounds["bottom"], b["bottom"]) + current_text += " " + txt + else: + # Save current block and start a new one + ocr_elements.append( + { + "text": current_text, + "bounds": ( + f"[{current_bounds['left']},{current_bounds['top']}][{current_bounds['right']},{current_bounds['bottom']}]" + ), + } + ) + current_bounds = b.copy() + current_text = txt + + # Don't forget the last block in the line + if current_bounds: + ocr_elements.append( + { + "text": current_text, + "bounds": ( + f"[{current_bounds['left']},{current_bounds['top']}][{current_bounds['right']},{current_bounds['bottom']}]" + ), + } + ) + + node["ocr_elements"] = ocr_elements + orig_class = str(node.get("className") or node.get("class") or "") + if "[OCR]" not in orig_class: + node["class"] = f"{orig_class} [OCR]".strip() + + return xml_hierarchy diff --git a/artemis/utils/python_executor.py b/artemis/utils/python_executor.py new file mode 100644 index 0000000..845c9d1 --- /dev/null +++ b/artemis/utils/python_executor.py @@ -0,0 +1,105 @@ +# 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 ast +import logging +from pathlib import Path +import queue +import time +import jupyter_client + +logger = logging.getLogger(__name__) + + +def validate_coder_script(code_str: str): + """Validates that the generated python code does not violate constraints. + + Specifically checks for usage of cv2.resize and numpy slicing which could + mess up the transformation tracking. + """ + try: + tree = ast.parse(code_str) + except SyntaxError as e: + raise ValueError(f"Syntax error in code: {e}") + + for node in ast.walk(tree): + # Prevent cv2.resize + if isinstance(node, ast.Attribute): + if node.attr == "resize": + # Quick check if it might be cv2.resize + raise ValueError( + "SECURITY/VALIDATION ERROR: Direct usage of resize() is" + " forbidden. You MUST use canvas.resize_by_factor(factor)." + ) + + # Prevent slicing (very aggressive, but forces them to use canvas.crop) + # Note: We cannot easily block all slicing because they might slice other arrays. + # But we can look for specific variable names or just warn. + # Let's be lenient on slicing but strict in the prompt. + + return True + + +class PythonExecutor: + def __init__(self, session_dir: Path | None = None): + self.session_dir = session_dir + try: + self.km = jupyter_client.KernelManager(kernel_name="python3") + self.km.start_kernel() + self.kc = self.km.client() + self.kc.start_channels() + self.kc.wait_for_ready(timeout=10) + except Exception as e: + logger.error(f"Failed to start Jupyter kernel: {e}") + raise RuntimeError(f"Kernel startup failed: {e}") + + def execute(self, code: str) -> str: + # First validate code + try: + validate_coder_script(code) + except ValueError as ve: + return str(ve) + + start_time = time.time() + self.kc.execute(code) + output_text = [] + + while True: + try: + msg = self.kc.get_iopub_msg(timeout=15) + msg_type = msg["header"]["msg_type"] + + if msg_type == "stream": + output_text.append(msg["content"]["text"]) + elif msg_type == "error": + traceback = "\n".join(msg["content"]["traceback"]) + output_text.append(f"Traceback:\n{traceback}") + elif msg_type in ("display_data", "execute_result"): + data = msg["content"]["data"] + if "text/plain" in data: + output_text.append(data["text/plain"]) + elif msg_type == "status" and msg["content"]["execution_state"] == "idle": + break + except queue.Empty: + output_text.append("[Execution Timeout Error]") + break + + return "".join(output_text) + + def close(self): + try: + self.kc.stop_channels() + self.km.shutdown_kernel() + except Exception as e: + logger.error(f"Failed to cleanly shutdown kernel: {e}") diff --git a/artemis/utils/recorder.py b/artemis/utils/recorder.py new file mode 100644 index 0000000..70dbee1 --- /dev/null +++ b/artemis/utils/recorder.py @@ -0,0 +1,60 @@ +# 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 time + +from langchain_core.messages import BaseMessage +from artemis.context import ArtemisContext +from artemis.controllers.controller_factory import create_device_controller +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +async def record_interaction(ctx: ArtemisContext, response: BaseMessage): + if not ctx.execution_setup: + raise ValueError("No execution setup found") + if not ctx.execution_setup.traces_path or not ctx.execution_setup.trace_name: + raise ValueError("No traces path or trace name found") + + logger.info("Recording interaction") + controller = create_device_controller(ctx) + screenshot_base64 = await controller.screenshot() + logger.info("Screenshot taken") + try: + controller = create_device_controller(ctx) + compressed_screenshot_base64 = controller.get_compressed_b64_screenshot(screenshot_base64) + except Exception as e: + logger.error(f"Error compressing screenshot: {e}") + return "Could not record this interaction" + timestamp = time.time() + folder = ctx.execution_setup.traces_path.joinpath(ctx.execution_setup.trace_name).resolve() + folder.mkdir(parents=True, exist_ok=True) + try: + with open( + folder.joinpath(f"{int(timestamp)}.jpeg").resolve(), + "wb", + ) as f: + f.write(base64.b64decode(compressed_screenshot_base64)) + + with open( + folder.joinpath(f"{int(timestamp)}.json").resolve(), + "w", + encoding="utf-8", + ) as f: + f.write(response.model_dump_json()) + except Exception as e: + logger.error(f"Error recording interaction: {e}") + return "Screenshot recorded successfully" diff --git a/artemis/utils/requests_utils.py b/artemis/utils/requests_utils.py new file mode 100644 index 0000000..aa56194 --- /dev/null +++ b/artemis/utils/requests_utils.py @@ -0,0 +1,51 @@ +# 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. + +from artemis.utils.logger import get_logger +import requests + +logger = get_logger(__name__) + + +def curl_from_request(req: requests.PreparedRequest) -> str: + """Converts a requests.PreparedRequest object to a valid cURL command string.""" + command = ["curl", f"-X {req.method}"] + + for key, value in req.headers.items(): + command.append(f'-H "{key}: {value}"') + + if req.body: + body = req.body + if isinstance(body, bytes): + body = body.decode("utf-8") + # Escape single quotes in the body for shell safety + body = body.replace("'", "'\\''") + command.append(f"-d '{body}'") + + command.append(f"'{req.url}'") + + return " ".join(command) + + +def logging_hook(response, *args, **kwargs): + """Hook to log the request as a cURL command.""" + curl_command = curl_from_request(response.request) + logger.debug(f"\n--- cURL Command ---\n{curl_command}\n--------------------") + + +def get_session_with_curl_logging() -> requests.Session: + """Returns a requests.Session with cURL logging enabled.""" + session = requests.Session() + session.hooks["response"] = [logging_hook] + return session diff --git a/artemis/utils/shell_utils.py b/artemis/utils/shell_utils.py new file mode 100644 index 0000000..b3c27ad --- /dev/null +++ b/artemis/utils/shell_utils.py @@ -0,0 +1,34 @@ +# 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 subprocess + + +def run_shell_command_on_host(command: str) -> str: + """Helper to run a shell command on the host and return the output.""" + try: + result = subprocess.run( + command, + shell=True, + check=True, + capture_output=True, + text=True, + ) + return result.stdout.strip() + except subprocess.CalledProcessError as e: + # Log the error and stderr for better debugging + error_message = ( + f"Command '{command}' failed with exit code {e.returncode}.\nStderr: {e.stderr.strip()}" + ) + raise RuntimeError(error_message) from e diff --git a/artemis/utils/task_tree.py b/artemis/utils/task_tree.py new file mode 100644 index 0000000..30e0e32 --- /dev/null +++ b/artemis/utils/task_tree.py @@ -0,0 +1,1077 @@ +# 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 ast +import hashlib +import json +from pathlib import Path +import re +from typing import Any + + +def format_action_clean(action_obj) -> str: + if not action_obj: + return "No Action" + if isinstance(action_obj, list): + if not action_obj: + return "No Action" + action_obj = action_obj[0] + + if not isinstance(action_obj, dict): + return str(action_obj) + + act_type = action_obj.get("action") or action_obj.get("name") + if not act_type: + return json.dumps(action_obj, ensure_ascii=False) + + target_text = action_obj.get("target_text") or action_obj.get("text") + coords = action_obj.get("coordinates") or action_obj.get("target") + app_name = ( + action_obj.get("app_name") + or action_obj.get("package_name") + or action_obj.get("package") + or action_obj.get("package_or_bundle_id") + or action_obj.get("bundle_id") + ) + keycode = action_obj.get("keycode") or action_obj.get("key") + + if act_type in ("tap", "click", "long_press"): + label = f"'{target_text}'" if target_text else "element" + if act_type == "long_press": + duration = action_obj.get("duration") or action_obj.get("duration_ms") + duration_str = f" for {duration}ms" if duration else "" + return f"Long pressed {label} at {coords}{duration_str}" + else: + times = action_obj.get("times") or action_obj.get("click_times") or 1 + try: + times = int(times) + except (ValueError, TypeError): + times = 1 + + if times == 2: + return f"Double tapped {label} at {coords}" + elif times > 2: + delay = action_obj.get("delay_ms") or action_obj.get("delay") + delay_str = f" with {delay}ms delay" if delay else "" + return f"Tapped {label} at {coords} {times} times{delay_str}" + else: + return f"Tapped {label} at {coords}" + elif act_type in ("input_text", "focus_and_input_text"): + label = f"'{target_text}'" if target_text else "field" + text_val = action_obj.get("text") + clear_exist = action_obj.get("clear_exist") or action_obj.get("clear_before_input") + clear_str = " (without clearing)" if clear_exist is False else "" + return f"Inputted '{text_val}' into {label} at {coords}{clear_str}" + elif act_type == "swipe": + direction = action_obj.get("action_direction") or action_obj.get("direction") + duration = action_obj.get("duration") or action_obj.get("duration_ms") + duration_str = f" over {duration}ms" if duration else "" + if direction: + return f"Swiped {direction}{duration_str}" + + start_coords = action_obj.get("start_coordinates") + end_coords = action_obj.get("end_coordinates") + if (not start_coords or not end_coords) and "coordinates" in action_obj: + coords_list = action_obj.get("coordinates") + if isinstance(coords_list, list) and len(coords_list) == 4: + start_coords = coords_list[:2] + end_coords = coords_list[2:] + return f"Swiped from {start_coords} to {end_coords}{duration_str}" + elif act_type == "press_key": + return f"Pressed key '{keycode}'" + elif act_type == "launch_app": + return f"Launched app '{app_name}'" + elif act_type == "stop_app": + return f"Stopped app '{app_name}'" + elif act_type == "manage_app": + intent = action_obj.get("intent") + return f"Managed app '{app_name}' (action: {intent})" + elif act_type == "wait_for_delay": + delay_ms = ( + action_obj.get("delay_ms") or action_obj.get("time_in_ms") or action_obj.get("duration") + ) + if not delay_ms: + delay_sec = action_obj.get("delay_seconds") or action_obj.get("delay") + if delay_sec: + try: + delay_ms = int(float(delay_sec) * 1000) + except ValueError: + delay_ms = f"{delay_sec}s" + delay_val = ( + f"{delay_ms}ms" + if isinstance(delay_ms, int) or (isinstance(delay_ms, str) and delay_ms.isdigit()) + else str(delay_ms) + ) + return f"Waited for {delay_val}" + elif act_type in ("click_sequence", "tap_sequence"): + seq = action_obj.get("target") or action_obj.get("sequence") or [] + return f"Tapped sequence of targets: {seq}" + else: + return f"Action: {act_type} with args: {json.dumps(action_obj, ensure_ascii=False)}" + + +def format_result_clean(result_obj) -> str | None: + if not result_obj: + return None + if not isinstance(result_obj, dict): + return str(result_obj) + + status = result_obj.get("status") + repair_status = result_obj.get("repair_status") + + exec_list = result_obj.get("execution") or result_obj.get("executed_actions") + if isinstance(exec_list, list) and exec_list: + first_exec = exec_list[0] + if isinstance(first_exec, dict): + attempts = first_exec.get("attempts") + error = first_exec.get("error") or result_obj.get("error") + repair = first_exec.get("repair") + + if repair: + if repair_status == "fixed": + return f"Repaired: {repair}" + else: + return f"Error: {repair}" + elif attempts and status == "failed": + return f"Error: {' | '.join(attempts)}" + elif error: + return f"Error: {error}" + elif status == "error": + return "Error" + else: + return None + + error = result_obj.get("error") + if error: + return f"Error: {error}" + elif status == "error": + return "Error" + + return None + + +def safe_parse_validation_result(result: Any) -> list: + if isinstance(result, (list, tuple)): + return list(result) + if not isinstance(result, str): + return [] + + # Clean up enum references in string format + + cleaned = re.sub( + r"", + r"'\1'", + result.strip(), + ) + + # Try parsing with ast.literal_eval + + try: + parsed = ast.literal_eval(cleaned) + if isinstance(parsed, (list, tuple)): + return list(parsed) + except Exception: + pass + + # Fallback to json.loads if literal_eval failed + try: + parsed = json.loads(cleaned) + if isinstance(parsed, (list, tuple)): + return list(parsed) + except Exception: + pass + + return [] + + +def format_tool_call_clean(name, args, result) -> str | None: + ACTION_TOOLS = { + "click", + "click_sequence", + "tap_sequence", + "swipe", + "input_text", + "focus_and_input_text", + "press_key", + "long_press", + "launch_app", + "manage_app", + "stop_app", + "wait_for_delay", + "_exec_click", + "_exec_click_sequence", + "_exec_swipe", + "_exec_input_text", + "_exec_press_key", + "_exec_long_press", + "_exec_manage_app", + "_exec_wait_for_delay", + } + + # Custom rendering for safety net validations + if name in ("safety_net_validation", "safety_net_pixel_validation"): + res_list = safe_parse_validation_result(result) + if len(res_list) >= 3: + passed, category, detail = res_list[0], res_list[1], res_list[2] + if passed: + return "Passed" + else: + return f"Failed - {str(detail).strip()}" + return "Passed" if result else "Failed" + + # Custom rendering for Operator call to failure_analyzer agent + if name == "failure_analyzer": + clean_args = { + k: v + for k, v in args.items() + if k + not in ( + "state", + "tool_call_id", + "pre_screenshot", + "post_screenshot", + ) + } + args_str = json.dumps(clean_args, ensure_ascii=False) + if len(args_str) > 120: + args_str = args_str[:120] + "..." + + if isinstance(result, dict): + status = result.get("status", "unknown") + analysis = result.get("analysis") or result.get("reason") or "No analysis provided." + else: + status = args.get("status", "unknown") + analysis = args.get("analysis", "No analysis provided.") + if len(analysis) > 120: + analysis = analysis[:120].strip() + "..." + return f"`failure_analyzer({args_str})` -> Outcome: {status.upper()} (Analysis: {analysis})" + + # Custom rendering for Failure Analyzer final report + if name == "report_failure_analysis": + status = args.get("status", "unknown") + analysis = args.get("analysis", "No analysis provided.") + if len(analysis) > 120: + analysis = analysis[:120].strip() + "..." + return f"Outcome: {status.upper()} (Analysis: {analysis})" + + if name in ("read_note", "save_note", "update_note") and args.get("key") == "task_plan": + return None + + # Custom rendering for all Action Tools (including _exec_ tools) + if name in ACTION_TOOLS: + action_name = name[6:] if name.startswith("_exec_") else name + + # Convert arguments to action_obj format with thorough parameter mapping + action_obj = { + "action": action_name, + "intent": args.get("intent") or args.get("action"), + "target": (args.get("target") or args.get("coordinates") or args.get("sequence")), + "target_text": ( + args.get("text") or args.get("target_text") or args.get("target_text_or_desc") + ), + "text": args.get("text"), + "clear_exist": (args.get("clear_exist") or args.get("clear_before_input")), + "clear_before_input": args.get("clear_before_input"), + "keycode": args.get("keycode") or args.get("key"), + "key": args.get("key"), + "app_name": ( + args.get("app_name") + or args.get("package_name") + or args.get("package") + or args.get("package_or_bundle_id") + or args.get("bundle_id") + ), + "delay": (args.get("delay") or args.get("delay_ms") or args.get("time_in_ms")), + "delay_ms": args.get("delay_ms"), + "time_in_ms": args.get("time_in_ms"), + "duration": (args.get("duration") or args.get("duration_ms") or args.get("time_in_ms")), + "duration_ms": args.get("duration_ms"), + "times": args.get("times") or args.get("click_times"), + "click_times": args.get("click_times"), + "direction": args.get("direction") or args.get("action_direction"), + "action_direction": args.get("action_direction"), + } + + # Handle special FA swipe "action" parameter + if action_name == "swipe" and "action" in args: + action_val = args["action"] + if isinstance(action_val, list) and len(action_val) == 4: + action_obj["coordinates"] = action_val + elif isinstance(action_val, str): + action_obj["action_direction"] = action_val + + # Call format_action_clean to get factual natural language action + action_clean = format_action_clean(action_obj) + + # Determine if there was an error in the result + error_detail = None + if isinstance(result, dict): + error_detail = result.get("error") or result.get("message") + elif isinstance(result, (list, tuple)) and result: + first_val = result[0] + if isinstance(first_val, str) and any( + x in first_val.lower() for x in ("error", "failed", "timeout") + ): + error_detail = first_val + elif isinstance(result, str): + res_lower = result.lower() + if "error" in res_lower or "failed" in res_lower or "timeout" in res_lower: + error_detail = result + + if error_detail: + error_clean = str(error_detail).strip() + + # If the error is formatted as a Python tuple string, e.g. "('Error...', None, None, None)" + if error_clean.startswith("(") and error_clean.endswith(")"): + parts = error_clean[1:-1].split(",") + if parts: + first_part = parts[0].strip() + if (first_part.startswith("'") and first_part.endswith("'")) or ( + first_part.startswith('"') and first_part.endswith('"') + ): + first_part = first_part[1:-1] + error_clean = first_part + + for prefix in ( + "Error executing click: ", + "Error during click: ", + "Swipe failed: ", + "Error during swipe: ", + "Error executing key press: ", + "Error executing app launch: ", + "Error launching app: ", + ): + if error_clean.startswith(prefix): + error_clean = error_clean[len(prefix) :] + return f"{action_clean} -> (Execution failed: {error_clean})" + else: + return action_clean + + clean_args = {k: v for k, v in args.items() if k not in ("state", "tool_call_id")} + args_str = json.dumps(clean_args, ensure_ascii=False) + + # Clean the result + res_str = "" + if isinstance(result, str) and result.startswith("["): + try: + parsed = json.loads(result) + if isinstance(parsed, list): + result = parsed + except (json.JSONDecodeError, TypeError): + pass + + if isinstance(result, list): + text_blocks = [] + for block in result: + if isinstance(block, dict): + if block.get("type") == "text": + text_blocks.append(block.get("text", "").strip()) + elif block.get("type") == "image_url": + pass + if text_blocks: + res_str = "\n".join(text_blocks) + else: + res_str = str(result) + elif isinstance(result, dict): + error = result.get("error") + if error: + res_str = f"Error: {error}" + else: + res_str = json.dumps(result, ensure_ascii=False) + if len(res_str) > 120: + res_str = res_str[:120] + "..." + else: + res_str = str(result) + if len(res_str) > 120: + res_str = res_str[:120] + "..." + + return f"`{name}({args_str})` -> {res_str}" + + +def format_step_action_result(step: dict) -> str: + action = step.get("action_taken") + action_parsed = action[0] if isinstance(action, list) and action else action + action_clean = "No Action" + if action_parsed: + action_clean = format_action_clean(action_parsed) + + intercepted = False + safety_net_detail = None + interleaved = step.get("interleaved_events") or [] + for event in interleaved: + if event.get("type") == "tool_call" and event.get("name") in ( + "safety_net_validation", + "safety_net_pixel_validation", + ): + res_list = safe_parse_validation_result(event.get("result")) + if len(res_list) > 0 and not res_list[0]: + intercepted = True + if len(res_list) >= 3: + safety_net_detail = res_list[2] + break + + if not intercepted: + result_obj = step.get("last_execution_result") + if isinstance(result_obj, dict): + exec_list = result_obj.get("execution") or [] + if exec_list: + attempts = exec_list[0].get("attempts") or [] + if attempts and any( + "pre-execution validation" in str(att).lower() + or "validation failed" in str(att).lower() + for att in attempts + ): + intercepted = True + safety_net_detail = attempts[0] + + recovery_actions = [] + for event in interleaved: + e_type = event.get("type") + name = event.get("name") or "" + if ( + e_type == "tool_call" + and name.startswith("_exec_") + and name != "report_failure_analysis" + ): + formatted = format_tool_call_clean(name, event.get("args") or {}, event.get("result")) + if formatted: + recovery_actions.append(formatted) + + if not recovery_actions: + legacy_tool_calls = step.get("tool_calls") or [] + for tc in legacy_tool_calls: + name = tc.get("name") or "" + if name.startswith("_exec_") and name != "report_failure_analysis": + payload = tc.get("payload") or {} + tc_args = payload.get("args") or tc.get("args") or {} + tc_result = ( + payload.get("result") or payload.get("error") or tc.get("result") or "No result" + ) + formatted = format_tool_call_clean(name, tc_args, tc_result) + if formatted: + recovery_actions.append(formatted) + + output_parts = [action_clean] + + if intercepted: + detail_msg = f": {str(safety_net_detail).strip()}" if safety_net_detail else "" + output_parts.append(f"(Action not executed{detail_msg})") + + if not intercepted and not recovery_actions: + result_obj = step.get("last_execution_result") + if isinstance(result_obj, dict) and result_obj.get("status") == "failed": + error_msg = None + exec_list = result_obj.get("execution") or [] + if exec_list and isinstance(exec_list[0], dict): + error_msg = exec_list[0].get("error") + if not error_msg: + error_msg = result_obj.get("error") + if error_msg: + output_parts.append(f"(Execution failed: {str(error_msg).strip()})") + + if recovery_actions: + recovery_str = ", ".join(recovery_actions) + output_parts.append(f"Recovery Actions: [{recovery_str}]") + + return " -> ".join(output_parts) + + +def _render_step_detailed( + step: dict, + relative_time: str, + summary: str, + action: Any, + result: Any, + is_most_recent: bool, + for_failure_analyzer: bool = False, +) -> str: + # 1. Step Header + if is_most_recent and for_failure_analyzer: + status_str = ( + "Most Recent Step (Failed to execute, this is the step you need to" + " focus on and repair), " + ) + else: + status_str = "Most Recent Step, " if is_most_recent else "" + + # Summary is omitted from the detailed view header as it is already fully detailed below. + step_line = f"- **Step {step['step_number']} ({status_str}Start: {relative_time})**" + + interleaved = step.get("interleaved_events") or [] + + # 2. Determine if Safety Net validation failed (intercepted) + intercepted = False + safety_net_detail = None + for event in interleaved: + if event.get("type") == "tool_call" and event.get("name") in ( + "safety_net_validation", + "safety_net_pixel_validation", + ): + res_list = safe_parse_validation_result(event.get("result")) + if len(res_list) > 0 and not res_list[0]: + intercepted = True + if len(res_list) >= 3: + safety_net_detail = res_list[2] + break + + # Fallback check on result structure for legacy/non-interleaved steps + if not intercepted and result: + exec_list = result.get("execution") or [] + if exec_list: + attempts = exec_list[0].get("attempts") or [] + if attempts and any( + "pre-execution validation" in str(att).lower() + or "validation failed" in str(att).lower() + for att in attempts + ): + intercepted = True + safety_net_detail = attempts[0] + + # 3. Categorize events into Operator and Failure Analyzer + operator_events = [] + failure_analyzer_events = [] + + # Internal tools list to filter out from Operator + INTERNAL_SYSTEM_TOOLS = { + "safety_net_validation", + "safety_net_pixel_validation", + "failure_analyzer", + "report_failure_analysis", + "hopper", + } + + if interleaved: + for event in interleaved: + e_type = event.get("type") + name = event.get("name") or "" + content = event.get("content") or "" + + # Skip the failure_analyzer agent call itself to prevent duplicates + if name == "failure_analyzer": + continue + + if e_type in ("thought", "native_thought"): + operator_events.append(event) + elif ( + e_type == "tool_call" + and name not in INTERNAL_SYSTEM_TOOLS + and not name.startswith("_exec_") + ): + operator_events.append(event) + elif e_type in ( + "failure_analyzer_thought", + "failure_analyzer_native_thought", + ): + failure_analyzer_events.append(event) + elif e_type == "tool_call" and ( + name.startswith("_exec_") or name == "report_failure_analysis" + ): + failure_analyzer_events.append(event) + else: + # Fallback for legacy/non-interleaved steps + raw_thinking = step.get("operator_raw_thinking") + native_thinking = step.get("operator_native_thinking") + if native_thinking: + operator_events.append({"type": "native_thought", "content": native_thinking}) + if raw_thinking: + operator_events.append({"type": "thought", "content": raw_thinking}) + + legacy_tool_calls = step.get("tool_calls") or [] + for tc in legacy_tool_calls: + name = tc.get("name") + if name in INTERNAL_SYSTEM_TOOLS or name.startswith("_exec_"): + # If it's failure analyzer tools in legacy steps, place them in failure analyzer + if name.startswith("_exec_") or name == "report_failure_analysis": + # Convert to interleaved style format + payload = tc.get("payload") or {} + tc_args = payload.get("args") or tc.get("args") or {} + tc_result = ( + payload.get("result") + or payload.get("error") + or tc.get("result") + or "No result" + ) + failure_analyzer_events.append( + { + "type": "tool_call", + "name": name, + "args": tc_args, + "result": tc_result, + } + ) + else: + payload = tc.get("payload") or {} + tc_args = payload.get("args") or tc.get("args") or {} + tc_result = ( + payload.get("result") or payload.get("error") or tc.get("result") or "No result" + ) + operator_events.append( + { + "type": "tool_call", + "name": name, + "args": tc_args, + "result": tc_result, + } + ) + + # Helper to check if a thought content is already in operator_events + def thought_exists(content): + content_clean = content.strip().lower() + for ev in operator_events: + if ( + ev.get("type") in ("thought", "native_thought") + and ev.get("content", "").strip().lower() == content_clean + ): + return True + return False + + # Extract thought from final physical action if present + action_thought = None + if action: + action_parsed = action[0] if isinstance(action, list) and action else action + if isinstance(action_parsed, dict): + action_thought = action_parsed.get("thought") + if action_thought and action_thought.strip() and not thought_exists(action_thought): + operator_events.insert(0, {"type": "thought", "content": action_thought.strip()}) + + # Extract native thinking from LLM thinking block first if not already present + native_thinking = step.get("operator_native_thinking") + if native_thinking and native_thinking.strip() and not thought_exists(native_thinking): + operator_events.append({"type": "native_thought", "content": native_thinking.strip()}) + + # Extract raw thinking from LLM text content second if not already present + raw_thinking = step.get("operator_raw_thinking") + if raw_thinking and raw_thinking.strip() and not thought_exists(raw_thinking): + operator_events.append({"type": "thought", "content": raw_thinking.strip()}) + + # 4. Render Operator Loop + if operator_events and (not for_failure_analyzer or is_most_recent): + step_line += "\n * [Operator Decision Loop]:" + ACTION_TOOLS = { + "click", + "swipe", + "input_text", + "press_key", + "long_press", + "launch_app", + "wait_for_delay", + } + for event in operator_events: + e_type = event["type"] + content = event.get("content") or "" + name = event.get("name") or "" + if not content.strip() and e_type != "tool_call": + continue + if e_type in ("native_thought", "thought"): + cleaned = re.sub( + r".*?", + "", + content, + flags=re.DOTALL, + ) + cleaned = re.sub(r"", "", cleaned, flags=re.DOTALL).strip() + if cleaned: + step_line += f"\n - {cleaned}" + elif e_type == "tool_call": + if name in ACTION_TOOLS: + continue + formatted = format_tool_call_clean( + name, event.get("args") or {}, event.get("result") + ) + if formatted: + step_line += f"\n - [Tool Call]: {formatted}" + + # Check if the execution failed + exec_error = None + if result and isinstance(result, dict): + status = result.get("status") + if status == "failed" or status == "error": + exec_list = result.get("execution") or result.get("executed_actions") + if isinstance(exec_list, list) and exec_list and isinstance(exec_list[0], dict): + first_exec = exec_list[0] + attempts = first_exec.get("attempts") + if attempts: + exec_error = " | ".join(attempts) + else: + exec_error = first_exec.get("error") or result.get("error") + else: + exec_error = result.get("error") + + # 5. Render Planned Action + if action: + action_clean = format_action_clean(action) + if intercepted: + step_line += ( + f"\n * [Planned Action]: {action_clean} (Intercepted by Pre-Execution Safety Net)" + ) + else: + if exec_error: + step_line += ( + f"\n * [Planned Action]: {action_clean} -> (Execution failed: {exec_error})" + ) + else: + step_line += f"\n * [Planned Action]: {action_clean}" + + # 6. Render Pre-Execution Safety Net + if intercepted: + step_line += "\n * [Pre-Execution Safety Net]:" + detail_str = ( + f"(Action not executed: {str(safety_net_detail).strip()})" + if safety_net_detail + else "(Action not executed)" + ) + step_line += f"\n - [Safety Net Check]: {detail_str}" + + # 7. Render Failure Analyzer Recovery Loop + if failure_analyzer_events: + step_line += "\n * [Failure Analyzer Recovery Loop]:" + for event in failure_analyzer_events: + e_type = event["type"] + content = event.get("content") or "" + name = event.get("name") or "" + if not content.strip() and e_type != "tool_call": + continue + if e_type in ( + "failure_analyzer_native_thought", + "failure_analyzer_thought", + "failure_analyzer_monologue", + ): + cleaned = content.strip() + if cleaned: + step_line += f"\n - {cleaned}" + elif e_type == "tool_call": + if name == "report_failure_analysis": + continue + formatted = format_tool_call_clean( + name, event.get("args") or {}, event.get("result") + ) + if formatted: + step_line += f"\n - [Tool Call]: {formatted}" + + # Check if Failure Analyzer was active and reported + has_fa = bool(failure_analyzer_events) + + # We consider FA as reported if it called the report tool, + # OR if the step's last_execution_result already contains a legacy repair description! + has_repair_in_result = False + if result and isinstance(result, dict): + exec_list = result.get("execution") or result.get("executed_actions") + if isinstance(exec_list, list) and exec_list and isinstance(exec_list[0], dict): + if exec_list[0].get("repair"): + has_repair_in_result = True + + fa_reported = ( + any( + event.get("type") == "tool_call" and event.get("name") == "report_failure_analysis" + for event in failure_analyzer_events + ) + or has_repair_in_result + ) + + # 8. Validator Execution Result + if has_fa and not fa_reported: + step_line += ( + "\n * [Result]: Unknown repair outcome (Failure Analyzer did not" + " submit a final report)" + ) + elif result: + result_clean = format_result_clean(result) + if result_clean: + step_line += f"\n * [Result]: {result_clean}" + + return step_line + + +def build_plan_and_history( + task_plan: str, + steps: list, + current_subgoal_hash: str, + keep_subgoal_hashes: set | None = None, + min_summaries: int = 5, + last_n_detailed: int = 1, + detailed_subgoal_hashes: set | None = None, + all_detailed: bool = False, + strict_milestone_pruning: bool = False, + recent_window_size: int = 3, + chronological_last_step: bool = False, + for_failure_analyzer: bool = False, +) -> str: + """Builds a clean plan list and separate flat chronological execution history with adaptive compression.""" + # 1. Plan Checklist Section + output_parts = [] + + # Check if task_plan has actual subgoals (e.g. starts with '- [') + has_subgoals = task_plan and any( + line.strip().startswith("- [") for line in task_plan.split("\n") + ) + + if has_subgoals: + output_parts.append("--- Task Plan ---") + output_parts.append(task_plan) + if for_failure_analyzer: + output_parts.append( + "*(Note: Provided for context only. Do not execute the remaining plan.)*" + ) + output_parts.append("") + elif not for_failure_analyzer: + # For non-failure_analyzer context, print fallback if not empty + if task_plan and task_plan != "No task plan yet.": + output_parts.extend( + [ + "--- Task Plan ---", + task_plan, + "", + ] + ) + + output_parts.append("--- Execution History ---") + + # 2. Chronological Steps History Section + if not steps: + output_parts.append("No steps executed yet.") + else: + full_info_step_id = steps[-1]["step_id"] + visible_step_ids = set() + visible_step_ids.add(full_info_step_id) + + # Calculate which steps are visible using milestone filtering or sliding window compression + if keep_subgoal_hashes is not None: + # Milestone-specific filtering (checker / diagnoser) + for step in steps: + meta = step.get("extra_metadata") or {} + step_subgoal_hash = meta.get("subgoal_hash") + if step_subgoal_hash in keep_subgoal_hashes: + visible_step_ids.add(step["step_id"]) + else: + # Default sliding window compression or strict milestone pruning + if strict_milestone_pruning: + recent_step_ids = set() + actual_recent_n = min(len(steps), recent_window_size) + for step in steps[-actual_recent_n:]: + recent_step_ids.add(step["step_id"]) + + for i in range(len(steps) - 1, -1, -1): + step = steps[i] + step_id = step["step_id"] + if step_id == full_info_step_id: + continue + + meta = step.get("extra_metadata") or {} + step_subgoal_hash = meta.get("subgoal_hash") + + # Strict pruning: only keep if active subgoal OR in recent continuity window + if step_subgoal_hash == current_subgoal_hash or step_id in recent_step_ids: + visible_step_ids.add(step_id) + else: + summary_count = 0 + for i in range(len(steps) - 1, -1, -1): + step = steps[i] + step_id = step["step_id"] + if step_id == full_info_step_id: + continue + + meta = step.get("extra_metadata") or {} + step_subgoal_hash = meta.get("subgoal_hash") + + # Keep steps under the active milestone, or slide back until min_summaries is met + if step_subgoal_hash == current_subgoal_hash or summary_count < min_summaries: + visible_step_ids.add(step_id) + summary_count += 1 + + detailed_step_ids = set() + # Pre-calculate steps that MUST be rendered in detailed mode (last N detailed) + if last_n_detailed > 0: + actual_n = min(len(steps), last_n_detailed) + for step in steps[-actual_n:]: + detailed_step_ids.add(step["step_id"]) + + # Ensure the last step is always in detailed_step_ids + detailed_step_ids.add(full_info_step_id) + + for step in steps: + step_id = step["step_id"] + + # Skip steps that are compressed (not marked as visible) + if step_id not in visible_step_ids: + continue + + relative_time = step.get("relative_time") or "N/A" + summary = step.get("summary") + action = step.get("action_taken") + thinking = step.get("operator_raw_thinking") + result = step.get("last_execution_result") + + # Determine if this step should be detailed (e.g. current/recent step or missing summary) + render_as_detailed = ( + all_detailed + or (step_id in detailed_step_ids) + or (not summary or summary == "Action executed.") + ) + + is_most_recent = step_id == full_info_step_id + if render_as_detailed: + step_line = _render_step_detailed( + step, + relative_time, + summary, + action, + result, + is_most_recent, + for_failure_analyzer, + ) + else: + step_line = f"- *Step {step['step_number']} (Start: {relative_time}): {summary}*" + if action: + action_clean = format_step_action_result(step) + step_line += f"\n *Action*: {action_clean}" + + output_parts.append(step_line) + + return "\n".join(output_parts) + + +def get_active_subgoal_hashes(task_plan: str) -> tuple[str, str | None]: + """Parses task_plan content to find the active top-level subgoal hash. + + Consolidates all active subtasks under their parent Level 1 subgoal, + returning (parent_hash, None). Falls back to the first pending subgoal if no + active subgoal is found and all are pending. + """ + if not task_plan: + return "default", None + + try: + lines = task_plan.split("\n") + active_line_idx = -1 + # Search from bottom to top to prioritize the most specific deeply nested active subgoal + for i in range(len(lines) - 1, -1, -1): + if lines[i].strip().startswith("- [/]"): + active_line_idx = i + break + + if active_line_idx != -1: + active_line = lines[active_line_idx] + active_text = active_line.strip()[5:].strip() + + # Check if it is indented (sub-subgoal) + if active_line.startswith(" "): + # Find parent (first non-indented line above it) + parent_text = "" + for i in range(active_line_idx - 1, -1, -1): + if not lines[i].startswith(" ") and lines[i].strip().startswith("- ["): + parent_text = lines[i].strip()[5:].strip() + break + + # Consolidate under the parent Level 1 subgoal + if parent_text: + parent_hash = hashlib.md5(parent_text.encode("utf-8")).hexdigest() + return parent_hash, None + + # If not indented, it is a top-level Level 1 subgoal + active_hash = hashlib.md5(active_text.encode("utf-8")).hexdigest() + return active_hash, None + + # No active subgoal found. Check if all are completed. + top_level_lines = [line for line in lines if line.strip().startswith("- [")] + all_done = all(line.strip().startswith("- [x]") for line in top_level_lines) + if all_done: + return "default", None + + # SAFE FALLBACK: Check if all subgoals are not started (all are "[ ]") + all_pending = all(line.strip().startswith("- [ ]") for line in top_level_lines) + if all_pending: + # If all are pending, we are at the start of the task. + # The active subgoal is guaranteed to be the first one. + for line in lines: + stripped = line.strip() + if stripped.startswith("- [ ]"): + first_pending_text = stripped[5:].strip() + first_hash = hashlib.md5(first_pending_text.encode("utf-8")).hexdigest() + return first_hash, None + + except Exception: + pass + + return "default", None + + +def get_completed_subgoal_hashes(task_plan: str) -> set[str]: + """Extracts MD5 hashes of all completed [x] top-level subgoals from task_plan.""" + if not task_plan: + return set() + + completed_hashes = set() + + pattern = re.compile(r"^-\s*\[x\]\s*(.*)$", re.MULTILINE) + for match in pattern.findall(task_plan): + completed_hashes.add(hashlib.md5(match.strip().encode("utf-8")).hexdigest()) + + return completed_hashes + + +def get_all_subgoal_aliases(target_hash: str, base_dir: str | Path | None = None) -> set[str]: + """Transitively resolves all older subgoal hashes that map to target_hash in subgoal_hash_chain.json.""" + aliases = {target_hash} + if not base_dir or not isinstance(base_dir, (str, Path)): + return aliases + + chain_path = Path(base_dir) / "notes" / "subgoal_hash_chain.json" + if not chain_path.exists(): + return aliases + + try: + chain = json.loads(chain_path.read_text(encoding="utf-8")) + except Exception: + return aliases + + # Transitively resolve old_hash -> new_hash chains + for old_hash in chain: + curr = old_hash + visited = set() + while curr in chain and curr not in visited: + visited.add(curr) + curr = chain[curr] + if curr == target_hash: + aliases.add(old_hash) + + return aliases + + +def get_recent_subgoal_hashes( + steps: list, current_subgoal_hash: str, base_dir: str | Path | None = None +) -> set[str]: + """Helper to get a set of subgoal hashes containing the current active subgoal hash, + + its historical failed/renamed aliases (resolved via + subgoal_hash_chain.json), + and the most recent completed subgoal hash as transition context. + """ + # 1. Resolve all alias hashes for the active subgoal + aliases = get_all_subgoal_aliases(current_subgoal_hash, base_dir) + keep_hashes = set(aliases) + + # 2. Find the most recent completed subgoal in history + # (meaning a step whose hash is NOT an alias of the active subgoal) + recent_completed_hash = None + for step in reversed(steps): + meta = step.get("extra_metadata") or {} + s_hash = meta.get("subgoal_hash") + if s_hash and s_hash not in aliases: + recent_completed_hash = s_hash + break + + if recent_completed_hash: + keep_hashes.add(recent_completed_hash) + + return keep_hashes diff --git a/artemis/utils/text.py b/artemis/utils/text.py new file mode 100644 index 0000000..3050aa2 --- /dev/null +++ b/artemis/utils/text.py @@ -0,0 +1,37 @@ +# 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. + +"""Text extraction and manipulation utilities.""" + +from typing import Any + + +def safe_extract_text(content: Any) -> str: + """Safely extracts a flat string from potentially complex or nested LLM response content. + + Guarantees returning a string object to prevent AttributeError when calling + string methods. + """ + if isinstance(content, str): + return content + elif isinstance(content, list): + return "".join(safe_extract_text(item) for item in content) + elif isinstance(content, dict): + if "text" in content: + return safe_extract_text(content["text"]) + return str(content) + elif content is None: + return "" + else: + return str(content) diff --git a/artemis/utils/time.py b/artemis/utils/time.py new file mode 100644 index 0000000..5ec443a --- /dev/null +++ b/artemis/utils/time.py @@ -0,0 +1,20 @@ +# 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. + +from datetime import datetime + + +def convert_timestamp_to_str(ts: float) -> str: + dt = datetime.fromtimestamp(ts) + return dt.strftime("%Y-%m-%dT%H-%M-%S") diff --git a/artemis/utils/ui_filter.py b/artemis/utils/ui_filter.py new file mode 100644 index 0000000..6f9f5ef --- /dev/null +++ b/artemis/utils/ui_filter.py @@ -0,0 +1,646 @@ +# 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 re +from typing import Any + +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + + +def _parse_bounds(bounds: Any) -> dict[str, int] | None: + """Normalize bounds from different formats to a standard dictionary. + + Supports: 1. String format: "[x1,y1][x2,y2]" 2. Dictionary format: {"x": + int, "y": int, "width": int, "height": int} 3. Dictionary format: {"left": + int, "top": int, "right": int, "bottom": int} 4. Cached format: + "parsed_bounds" inside element dict + """ + if not bounds: + return None + + if isinstance(bounds, dict): + if "parsed_bounds" in bounds and isinstance(bounds["parsed_bounds"], dict): + return bounds["parsed_bounds"] + if "left" in bounds and "top" in bounds: + return { + "left": int(bounds.get("left", 0)), + "top": int(bounds.get("top", 0)), + "right": int(bounds.get("right", 0)), + "bottom": int(bounds.get("bottom", 0)), + } + if "x" in bounds and "y" in bounds: + x = int(bounds.get("x", 0)) + y = int(bounds.get("y", 0)) + width = int(bounds.get("width", 0)) + height = int(bounds.get("height", 0)) + return { + "left": x, + "top": y, + "right": x + width, + "bottom": y + height, + } + if "bounds" in bounds and isinstance(bounds["bounds"], str): + res = _parse_bounds(bounds["bounds"]) + if res: + bounds["parsed_bounds"] = res + return res + + if isinstance(bounds, str): + match = re.match(r"\[(\d+),(\d+)\]\[(\d+),(\d+)\]", bounds) + if match: + x1, y1, x2, y2 = map(int, match.groups()) + return {"left": x1, "top": y1, "right": x2, "bottom": y2} + return None + return None + + +def _intersects_screen(bounds: dict[str, int], screen_width: int, screen_height: int) -> bool: + """Check if element intersects screen bounds.""" + left = bounds.get("left", 0) + top = bounds.get("top", 0) + right = bounds.get("right", 0) + bottom = bounds.get("bottom", 0) + return not (right <= 0 or bottom <= 0 or left >= screen_width or top >= screen_height) + + +def _meets_min_size(bounds: dict[str, int], min_size: int) -> bool: + """Check if element meets minimum size requirements.""" + width = bounds.get("right", 0) - bounds.get("left", 0) + height = bounds.get("bottom", 0) - bounds.get("top", 0) + return width > min_size and height > min_size + + +def _clip_bounds( + node: dict[str, Any], + bounds_data: dict[str, int], + ancestor_bounds: dict[str, int], +) -> dict[str, int]: + """Clips bounds to ancestor bounds and updates the node inplace.""" + left = max(bounds_data["left"], ancestor_bounds.get("left", 0)) + top = max(bounds_data["top"], ancestor_bounds.get("top", 0)) + right = min(bounds_data["right"], ancestor_bounds.get("right", 0)) + bottom = min(bounds_data["bottom"], ancestor_bounds.get("bottom", 0)) + + is_fully_clipped = False + if left >= right or top >= bottom: + is_fully_clipped = True + # Clamp to a point on the edge of ancestor bounds + left = right = max( + ancestor_bounds.get("left", 0), + min(bounds_data["left"], ancestor_bounds.get("right", 0)), + ) + top = bottom = max( + ancestor_bounds.get("top", 0), + min(bounds_data["top"], ancestor_bounds.get("bottom", 0)), + ) + + new_bounds = {"left": left, "top": top, "right": right, "bottom": bottom} + + if is_fully_clipped: + node["is_clipped"] = True + + # Re-serialize bounds into the node matching original format + orig = node.get("bounds") + if isinstance(orig, str): + node["bounds"] = f"[{left},{top}][{right},{bottom}]" + elif isinstance(orig, dict): + if "left" in orig and "top" in orig: + node["bounds"] = new_bounds.copy() + elif "x" in orig and "y" in orig: + node["bounds"] = { + "x": left, + "y": top, + "width": right - left, + "height": bottom - top, + } + return new_bounds + + +def _is_semantic_empty(node: dict[str, Any]) -> bool: + """Check if a node has no semantic value or interactive capability.""" + # Check text + text = node.get("text") + if text and str(text).strip(): + return False + + # Check content description + content_desc = ( + node.get("content-desc") + or node.get("contentDescription") + or node.get("content-description") + ) + if content_desc and str(content_desc).strip(): + return False + + # Check clickable + clickable = node.get("clickable") + if clickable is True or str(clickable).lower() == "true": + return False + + # Check focusable + focusable = node.get("focusable") + if focusable is True or str(focusable).lower() == "true": + return False + + # Check if it's a valid image container + class_name = str(node.get("className") or node.get("class") or "").lower() + if "image" in class_name or "icon" in class_name or "photo" in class_name: + return False + + return True + + +def _merge_complementary_children(children: list[dict], parent_node: dict) -> list[dict]: + """Strategy 2: Merge complementary siblings under a clickable parent.""" + # Rule 1: Parent must be clickable or focusable + clickable = parent_node.get("clickable") + focusable = parent_node.get("focusable") + if not ( + clickable is True + or str(clickable).lower() == "true" + or focusable is True + or str(focusable).lower() == "true" + ): + return children + + # Classify children + textual_children = [] + decoration_children = [] + + for child in children: + # Rule 3: No child is a form control + class_name = str(child.get("className") or child.get("class") or "").lower() + if any(kw in class_name for kw in ["checkbox", "switch", "radio", "edittext"]): + return children + + text = child.get("text") + content_desc = ( + child.get("content-desc") + or child.get("contentDescription") + or child.get("content-description") + ) + + if (text and str(text).strip()) or (content_desc and str(content_desc).strip()): + textual_children.append(child) + else: + # Check if it's a decoration (not clickable, not focusable) + c_clickable = child.get("clickable") + c_focusable = child.get("focusable") + if not ( + c_clickable is True + or str(c_clickable).lower() == "true" + or c_focusable is True + or str(c_focusable).lower() == "true" + ): + decoration_children.append(child) + else: + # It's an interactive element without text, don't merge as decoration + return children + + # Rule 2: Exactly ONE textual child and at least one decoration child + if ( + len(textual_children) == 1 + and len(decoration_children) >= 1 + and len(children) == len(textual_children) + len(decoration_children) + ): + # Perform merge + txt_child = textual_children[0] + + # Compute Union Bounds + all_bounds = [] + for c in children: + b = _parse_bounds(c.get("bounds")) + if b: + all_bounds.append(b) + + if not all_bounds: + return children + + left = min(b["left"] for b in all_bounds) + top = min(b["top"] for b in all_bounds) + right = max(b["right"] for b in all_bounds) + bottom = max(b["bottom"] for b in all_bounds) + + merged_node = txt_child.copy() + + # Update bounds matching original format if possible, or just dict + orig = txt_child.get("bounds") + if isinstance(orig, str): + merged_node["bounds"] = f"[{left},{top}][{right},{bottom}]" + else: + merged_node["bounds"] = { + "left": left, + "top": top, + "right": right, + "bottom": bottom, + } + + # Concatenate text and content-desc + desc_parts = [] + for c in children: + desc = ( + c.get("content-desc") or c.get("contentDescription") or c.get("content-description") + ) + if desc and str(desc).strip(): + desc_parts.append(f"[{str(desc).strip()}]") + elif "image" in str(c.get("className") or c.get("class") or "").lower(): + # User tip: use [Icon] for images without desc + desc_parts.append("[Icon]") + + text_parts = [] + for c in children: + txt = c.get("text") + if txt and str(txt).strip(): + text_parts.append(str(txt).strip()) + + if desc_parts: + merged_node["content-desc"] = " ".join(desc_parts) + if text_parts: + merged_node["text"] = " ".join(text_parts) + + return [merged_node] + + return children + + +def _prune_parent_child_redundancy( + parent: dict[str, Any], children: list[dict[str, Any]] +) -> list[dict[str, Any]]: + """Strategy 3: Remove redundant children that have the same text as the parent.""" + parent_text = str(parent.get("text") or "").strip() + parent_desc = str( + parent.get("content-desc") + or parent.get("contentDescription") + or parent.get("content-description") + or "" + ).strip() + + if not parent_text and not parent_desc: + return children + + new_children = [] + for child in children: + child_text = str(child.get("text") or "").strip() + child_desc = str( + child.get("content-desc") + or child.get("contentDescription") + or child.get("content-description") + or "" + ).strip() + + text_matches = False + if parent_text and child_text and parent_text == child_text: + text_matches = True + elif parent_desc and child_text and parent_desc == child_text: + text_matches = True + elif parent_text and child_desc and parent_text == child_desc: + text_matches = True + elif parent_desc and child_desc and parent_desc == child_desc: + text_matches = True + + if text_matches: + class_name = str(child.get("className") or child.get("class") or "").lower() + is_form_control = any( + kw in class_name for kw in ["checkbox", "switch", "radio", "edittext"] + ) + has_children = bool(child.get("children")) + + clickable = child.get("clickable") + focusable = child.get("focusable") + is_interactive = ( + clickable is True + or str(clickable).lower() == "true" + or focusable is True + or str(focusable).lower() == "true" + ) + + if is_form_control or has_children or is_interactive: + # Keep but erase text + merged_node = child.copy() + if "text" in merged_node: + merged_node["text"] = "" + if "content-desc" in merged_node: + merged_node["content-desc"] = "" + if "contentDescription" in merged_node: + merged_node["contentDescription"] = "" + if "content-description" in merged_node: + merged_node["content-description"] = "" + new_children.append(merged_node) + else: + # Remove child + continue + else: + new_children.append(child) + + return new_children + + +def _filter_and_prune( + node: dict[str, Any], + screen_width: int, + screen_height: int, + min_element_size: int = 5, + use_dynamic_min_size: bool = True, + clip_bounds: bool = True, + semantic_pruning: str = "hoist", + ancestor_bounds: dict[str, int] | None = None, +) -> list[dict[str, Any]]: + """Recursively filters, clips bounds, and applies semantic pruning strategies.""" + if ancestor_bounds is None: + ancestor_bounds = { + "left": 0, + "top": 0, + "right": screen_width, + "bottom": screen_height, + } + + bounds_data = _parse_bounds(node.get("bounds")) + + if bounds_data: + # We no longer check _intersects_screen to return [] immediately, + # to ensure zero info loss. Out-of-bounds elements will be clipped to a point. + + if clip_bounds: + bounds_data = _clip_bounds(node, bounds_data, ancestor_bounds) + + # Dynamic min size logic: use relative percentage (0.5% of short side) if enabled + actual_min_size = min_element_size + if use_dynamic_min_size: + short_side = min(screen_width, screen_height) + dynamic_min_size = int(short_side * 0.005) + actual_min_size = max(min_element_size, dynamic_min_size) + + # If it's fully clipped, we skip min size check to preserve it + if not node.get("is_clipped", False) and not _meets_min_size(bounds_data, actual_min_size): + return [] + + # Recursively process children + filtered_node = node.copy() + filtered_children = [] + if "children" in node: + for child in node.get("children", []): + if isinstance(child, dict): + filtered_children.extend( + _filter_and_prune( + child, + screen_width, + screen_height, + min_element_size, + use_dynamic_min_size, + clip_bounds, + semantic_pruning, + ancestor_bounds=bounds_data if bounds_data else ancestor_bounds, + ) + ) + + # Apply Strategy 3: Parent-Child Redundancy Removal + if filtered_children: + filtered_children = _prune_parent_child_redundancy(filtered_node, filtered_children) + + # Apply Strategy 2: Sibling Merging + if filtered_children: + filtered_children = _merge_complementary_children(filtered_children, filtered_node) + + filtered_node["children"] = filtered_children + + # Apply semantic pruning strategies + if _is_semantic_empty(filtered_node): + if semantic_pruning == "safe": + # Only eliminate if it's a leaf node + if not filtered_node.get("children"): + return [] + elif semantic_pruning == "hoist": + # Skip this node and hoist its valid children up + return filtered_node.get("children", []) + + return [filtered_node] + + +def _detect_fixed_system_bars( + ui_hierarchy: list[dict[str, Any]], screen_width: int, screen_height: int +) -> list[dict[str, int]]: + """Detect fixed, non-scrollable system or navigation bars anchored to the extreme top/bottom edges. + + Enforces strict triple guardrails to guarantee near-zero false positives: + + 1. Edge & Width Guarantee: Must anchor to top (top==0) or bottom + (bottom==screen_height) and span >= 95% of screen width. + 2. Non-Scrollable Root Branch: The candidate node and all its + descendants/ancestors must NOT be scrollable. + 3. Identity & Height Bounds: High-confidence system package/IDs OR strict + navigation bar height bounds. + """ + fixed_bars = [] + + def check_scrollable(node: dict[str, Any]) -> bool: + if node.get("scrollable") is True or str(node.get("scrollable", "")).lower() == "true": + return True + for child in node.get("children", []): + if isinstance(child, dict) and check_scrollable(child): + return True + return False + + def scan_for_bars(node: dict[str, Any], has_scrollable_ancestor: bool): + is_curr_scrollable = has_scrollable_ancestor or ( + node.get("scrollable") is True or str(node.get("scrollable", "")).lower() == "true" + ) + + bounds = _parse_bounds(node.get("bounds")) + if bounds and not is_curr_scrollable and not check_scrollable(node): + width = bounds["right"] - bounds["left"] + height = bounds["bottom"] - bounds["top"] + package_name = str(node.get("package") or node.get("packageName") or "").lower() + res_id = str(node.get("resource-id") or node.get("resourceId") or "").lower() + class_name = str(node.get("className") or node.get("class") or "").lower() + + is_full_width = width >= int(screen_width * 0.95) + is_high_conf_id = ( + any( + kw in res_id or kw in class_name + for kw in [ + "bottom_navigation", + "navigation_bar", + "bottom_bar", + "tab_layout", + "action_bar", + ] + ) + or package_name == "com.android.systemui" + ) + + # Check bottom fixed bar (e.g., bottom navigation bar / system navigation bar) + if bounds["bottom"] >= screen_height and is_full_width and height > 0: + if is_high_conf_id or (25 <= height <= int(screen_height * 0.15)): + fixed_bars.append(bounds.copy()) + return # Don't recurse into children of confirmed fixed bars + + # Check top fixed bar (e.g., status bar / top action bar) + if bounds["top"] <= 0 and is_full_width and height > 0: + if package_name == "com.android.systemui" or ( + 15 <= height <= int(screen_height * 0.12) + ): + fixed_bars.append(bounds.copy()) + return + + for child in node.get("children", []): + if isinstance(child, dict): + scan_for_bars(child, is_curr_scrollable) + + for root_node in ui_hierarchy: + if isinstance(root_node, dict): + scan_for_bars(root_node, False) + + return fixed_bars + + +def _clamp_by_fixed_bars( + node: dict[str, Any], fixed_bars: list[dict[str, int]], min_size: int +) -> bool: + """Non-destructively clamp the vertical bounds of a node if it overlaps with confirmed fixed system bars. + + Returns False if the remaining visible height is smaller than min_size + (silent discard of thin slivers). + """ + bounds = _parse_bounds(node.get("bounds")) + if not bounds or not fixed_bars: + return True + + left, top, right, bottom = ( + bounds["left"], + bounds["top"], + bounds["right"], + bounds["bottom"], + ) + + for bar in fixed_bars: + # Check horizontal overlap + if right <= bar["left"] or left >= bar["right"]: + continue + + # If it's the fixed bar itself or fully inside the bar, keep it if it belongs to the bar + if ( + left == bar["left"] + and top == bar["top"] + and right == bar["right"] + and bottom == bar["bottom"] + ): + return True + + # Clamp against a bottom fixed bar + if bar["bottom"] >= bottom and top < bar["top"] < bottom: + bottom = min(bottom, bar["top"]) + + # Clamp against a top fixed bar + if bar["top"] <= top and top < bar["bottom"] < bottom: + top = max(top, bar["bottom"]) + + # If remaining height is too small, silently discard this occluded sliver + if bottom - top <= max(min_size, 20): + return False + + new_bounds = {"left": left, "top": top, "right": right, "bottom": bottom} + orig = node.get("bounds") + if isinstance(orig, str): + node["bounds"] = f"[{left},{top}][{right},{bottom}]" + elif isinstance(orig, dict): + if "left" in orig and "top" in orig: + node["bounds"] = new_bounds.copy() + elif "x" in orig and "y" in orig: + node["bounds"] = { + "x": left, + "y": top, + "width": right - left, + "height": bottom - top, + } + return True + + +def filter_node( + node: dict[str, Any], + screen_width: int, + screen_height: int, + min_element_size: int = 5, + use_dynamic_min_size: bool = True, + clip_bounds: bool = True, + semantic_pruning: str = "hoist", +) -> dict[str, Any] | None: + """Recursively filter a UI node and its children based on screen bounds, size, and semantics. + + Returns None if the node does not meet the criteria. + """ + res_list = _filter_and_prune( + node, + screen_width, + screen_height, + min_element_size, + use_dynamic_min_size, + clip_bounds, + semantic_pruning, + ) + if not res_list: + return None + + if len(res_list) == 1: + return res_list[0] + + # If the root node itself was hoisted resulting in multiple children, wrap or attach + new_node = node.copy() + new_node["children"] = res_list + return new_node + + +def filter_ui_hierarchy( + ui_hierarchy: list[dict[str, Any]], + screen_width: int = 1080, + screen_height: int = 2400, + min_element_size: int = 5, + use_dynamic_min_size: bool = True, + clip_bounds: bool = True, + semantic_pruning: str = "hoist", +) -> list[dict[str, Any]]: + """Pure function to filter an entire list of UI elements (either flat or nested). + + Removes elements outside screen bounds, smaller than minimum size, or + semantically empty. Applies non-destructive safe-bounds clamping against + confirmed fixed system/navigation bars. + """ + # 1. Detect high-confidence fixed system bars before flattening + fixed_bars = _detect_fixed_system_bars(ui_hierarchy, screen_width, screen_height) + + filtered_elements = [] + for element in ui_hierarchy: + if isinstance(element, dict): + filtered_elements.extend( + _filter_and_prune( + element, + screen_width=screen_width, + screen_height=screen_height, + min_element_size=min_element_size, + use_dynamic_min_size=use_dynamic_min_size, + clip_bounds=clip_bounds, + semantic_pruning=semantic_pruning, + ) + ) + + # 2. Apply safe-bounds clamping against fixed system bars on final candidate elements + if fixed_bars: + clamped_elements = [] + for el in filtered_elements: + if _clamp_by_fixed_bars(el, fixed_bars, min_element_size): + clamped_elements.append(el) + return clamped_elements + + return filtered_elements diff --git a/artemis/utils/ui_hierarchy.py b/artemis/utils/ui_hierarchy.py new file mode 100644 index 0000000..6114497 --- /dev/null +++ b/artemis/utils/ui_hierarchy.py @@ -0,0 +1,166 @@ +# 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. + +from artemis.utils.logger import get_logger +from pydantic import BaseModel, Field + +logger = get_logger(__name__) + + +def __find_element_by_resource_id_in_rich_hierarchy( + hierarchy: list[dict], resource_id: str +) -> dict | None: + """Retrieves all the sibling elements for a given resource ID from a nested dictionary. + + Args: + hierarchy (dict): The nested dictionary representing the UI hierarchy. + resource_id (str): The resource-id to find. + + Returns: + list: A list of the sibling elements, or None if the resource_id is not + found. + """ + if not hierarchy: + return None + + for child in hierarchy: + if child.get("attributes", {}).get("resource-id") == resource_id: + return child.get("attributes", {}) + + for child in hierarchy: + result = __find_element_by_resource_id_in_rich_hierarchy( + child.get("children", []), resource_id + ) + if result is not None: + return result + + return None + + +def text_input_is_empty(text: str | None, hint_text: str | None) -> bool: + return not text or text == hint_text + + +def find_element_by_resource_id( + ui_hierarchy: list[dict], + resource_id: str, + index: int | None = None, + is_rich_hierarchy: bool = False, +) -> dict | None: + """Find a UI element by its resource-id in the UI hierarchy. + + Args: + ui_hierarchy: List of UI element dictionaries + resource_id: The resource-id to search for (e.g., + "com.google.android.settings.intelligence:id/open_search_view_edit_text") + + Returns: + The complete UI element dictionary if found, None otherwise + """ + if is_rich_hierarchy: + return __find_element_by_resource_id_in_rich_hierarchy(ui_hierarchy, resource_id) + + target_idx = index or 0 + + def search_recursive(elements: list[dict]) -> dict | None: + nonlocal target_idx + for element in elements: + if isinstance(element, dict): + if element.get("resourceId") == resource_id: + if target_idx == 0: + return element + target_idx -= 1 + continue + + children = element.get("children", []) + if children: + result = search_recursive(children) + if result: + return result + return None + + return search_recursive(ui_hierarchy) + + +def is_element_focused(element: dict) -> bool: + return element.get("focused", None) == "true" + + +def get_element_text(element: dict, hint_text: bool = False) -> str | None: + if hint_text: + return element.get("hintText", None) + return element.get("text", None) + + +class _CyFunctionDetectorMeta(type): + def __instancecheck__(self, instance): + name = type(instance).__name__ + return ( + name + in ( + "cyfunction", + "cython_function_or_method", + "builtin_function_or_method", + ) + or "cyfunction" in name.lower() + ) + + +class CyFunctionDetector(metaclass=_CyFunctionDetectorMeta): + pass + + +class Point(BaseModel): + model_config = {"ignored_types": (CyFunctionDetector,)} + x: int + y: int + + +class ElementBounds(BaseModel): + model_config = {"ignored_types": (CyFunctionDetector,)} + x: int = Field(description="The x coordinate of the top-left corner of the element.") + y: int = Field(description="The y coordinate of the top-left corner of the element.") + width: int = Field(description="The width of the element.") + height: int = Field(description="The height of the element.") + + def get_center(self) -> Point: + return Point(x=self.x + self.width // 2, y=self.y + self.height // 2) + + def get_relative_point(self, x_percent: float, y_percent: float) -> Point: + """Returns the coordinates of the point at x_percent of the width and y_percent + + of the height of the element. + + Ex if x_percent = 0.95 and y_percent = 0.95, + the point is at the bottom right of the element: + <------> + | | + | x| + <------> + """ + return Point( + x=int(self.x + self.width * x_percent), + y=int(self.y + self.height * y_percent), + ) + + +def get_bounds_for_element(element: dict) -> ElementBounds | None: + bounds = element.get("bounds", None) + if bounds: + try: + return ElementBounds(**bounds) + except Exception as e: + logger.error(f"Failed to validate bounds: {e}") + return None + return None diff --git a/artemis/utils/verification.py b/artemis/utils/verification.py new file mode 100644 index 0000000..3b50718 --- /dev/null +++ b/artemis/utils/verification.py @@ -0,0 +1,62 @@ +# 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 json +from pathlib import Path +from artemis.utils.logger import get_logger +from artemis.utils.notes import get_notes_dir + +logger = get_logger(__name__) + + +def get_verification_chat_path(base_dir: str | Path, subgoal_hash: str) -> Path: + """Gets the path to the verification chat file for a given subgoal.""" + notes_dir = get_notes_dir(base_dir) + return notes_dir / f"verification_chat_{subgoal_hash}.json" + + +def read_verification_chat(chat_path: Path) -> list[dict]: + """Reads the verification chat history from the given path.""" + if not chat_path.exists(): + return [] + try: + return json.loads(chat_path.read_text(encoding="utf-8")) + except Exception as e: + logger.error(f"Failed to load verification chat from {chat_path}: {e}") + return [] + + +def get_verification_chat_rounds(turns: list[dict]) -> tuple[int, int]: + """Calculates the maximum rounds for operator and checker from the turns.""" + max_op = max( + [t.get("round", 0) for t in turns if t.get("role") == "operator"], + default=0, + ) + max_chk = max( + [t.get("round", 0) for t in turns if t.get("role") == "checker"], + default=0, + ) + return max_op, max_chk + + +def append_verification_chat(chat_path: Path, role: str, content: str, round_num: int) -> bool: + """Appends a new turn to the verification chat file.""" + turns = read_verification_chat(chat_path) + turns.append({"role": role, "round": round_num, "content": content}) + try: + chat_path.write_text(json.dumps(turns, indent=2, ensure_ascii=False), encoding="utf-8") + return True + except Exception as e: + logger.error(f"Failed to save verification chat to {chat_path}: {e}") + return False diff --git a/artemis/utils/video.py b/artemis/utils/video.py new file mode 100644 index 0000000..17b85e2 --- /dev/null +++ b/artemis/utils/video.py @@ -0,0 +1,536 @@ +# 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. + +"""Video recording utilities for mobile devices. + +Provides shared types and utilities for video recording across platforms. +""" + +import asyncio +import importlib.util +from pathlib import Path +import platform +import shutil +import subprocess +from typing import Any +from uuid import UUID + +import cv2 +from pydantic import BaseModel, ConfigDict + +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) + +DEFAULT_MAX_DURATION_SECONDS = 900 # 15 minutes +VIDEO_READY_DELAY_SECONDS = 1 +ANDROID_DEVICE_VIDEO_PATH = "/sdcard/screen_recording.mp4" +ANDROID_MAX_RECORDING_DURATION_SECONDS = 180 # Android screenrecord limit + +# Expanded for Gemini File API (Supports up to 2GB). +# Target 100MB to allow 3-5min crisp video and prevent blurring for long durations. +MAX_VIDEO_SIZE_MB = 500 +MAX_VIDEO_SIZE_BYTES = MAX_VIDEO_SIZE_MB * 1024 * 1024 + + +class _CyFunctionDetectorMeta(type): + def __instancecheck__(self, instance): + name = type(instance).__name__ + return ( + name + in ( + "cyfunction", + "cython_function_or_method", + "builtin_function_or_method", + ) + or "cyfunction" in name.lower() + ) + + +class CyFunctionDetector(metaclass=_CyFunctionDetectorMeta): + pass + + +class RecordingSession(BaseModel): + """Tracks an active video recording session.""" + + model_config = ConfigDict(arbitrary_types_allowed=True) + + video_id: UUID + device_id: str + start_time: float + process: Any = None + data_engine_start_time: float | None = None + local_video_path: Path | None = None + android_device_path: str = ANDROID_DEVICE_VIDEO_PATH + android_video_segments: list[Path] = [] + android_segment_index: int = 0 + android_restart_task: asyncio.Task | None = None + watchdog_task: asyncio.Task | None = None + is_active: bool = True + errors: list[str] = [] + + +class VideoRecordingResult(BaseModel): + """Result of a video recording operation.""" + + success: bool + message: str + video_path: Path | None = None + file_size_mb: float | None = None + duration_seconds: float | None = None + actual_start_relative_time: float | None = None + warning: str | None = None + + model_config = ConfigDict(arbitrary_types_allowed=True) + + +# Global session storage - keyed by device_id +_active_recordings: dict[str, RecordingSession] = {} + + +def get_active_session(device_id: str) -> RecordingSession | None: + """Get the active recording session for a device.""" + return _active_recordings.get(device_id) + + +def set_active_session(device_id: str, session: RecordingSession) -> None: + """Set the active recording session for a device.""" + _active_recordings[device_id] = session + + +def remove_active_session(device_id: str) -> RecordingSession | None: + """Remove and return the active recording session for a device.""" + return _active_recordings.pop(device_id, None) + + +def has_active_session(device_id: str) -> bool: + """Check if there's an active recording session for a device.""" + return device_id in _active_recordings + + +def is_ffmpeg_installed() -> bool: + """Check if ffmpeg is available via imageio_ffmpeg or system PATH.""" + + if importlib.util.find_spec("imageio_ffmpeg") is not None: + return True + return shutil.which("ffmpeg") is not None + + +def get_ffmpeg_path() -> str: + """Get the path to the ffmpeg executable.""" + try: + import imageio_ffmpeg + + return imageio_ffmpeg.get_ffmpeg_exe() + except ImportError: + return "ffmpeg" + + +def is_scrcpy_installed() -> bool: + """Check if scrcpy is available in the system PATH.""" + + return shutil.which("scrcpy") is not None + + +def detect_video_tools_enabled() -> bool: + """Check if both scrcpy and ffmpeg are available to enable automated video features.""" + return is_ffmpeg_installed() and is_scrcpy_installed() + + +class FFmpegNotInstalledError(Exception): + """Raised when ffmpeg is required but not installed.""" + + def __init__(self): + os_name = platform.system().lower() + if os_name == "darwin": # macOS + install_instructions = "brew install ffmpeg" + elif os_name == "windows": + install_instructions = "Download from https://www.ffmpeg.org/download.html" + else: # Linux and others + install_instructions = ( + "Install via your package manager (e.g., apt install ffmpeg," + " dnf install ffmpeg) or download from" + " https://www.ffmpeg.org/download.html" + ) + + message = ( + "\n\nโŒ ffmpeg is required for video recording but is not" + " installed.\n\nPlease install ffmpeg first:\n โ†’" + f" {install_instructions}\n\nAfter installation, restart Artemis.\n" + ) + super().__init__(message) + + +def check_ffmpeg_available() -> None: + """Check if ffmpeg is installed and raise an error if not. + + Raises: + FFmpegNotInstalledError: If ffmpeg is not found in PATH. + """ + if not is_ffmpeg_installed(): + raise FFmpegNotInstalledError() + + +async def concatenate_videos(segments: list[Path], output_path: Path) -> bool: + """Concatenate multiple video segments using ffmpeg.""" + if not segments: + return False + + if len(segments) == 1: + shutil.move(segments[0], output_path) + return True + + list_file = output_path.parent / "segments.txt" + with open(list_file, "w") as f: + for segment in segments: + f.write(f"file '{segment}'\n") + + try: + process = await asyncio.create_subprocess_exec( + get_ffmpeg_path(), + "-y", + "-f", + "concat", + "-safe", + "0", + "-i", + str(list_file), + "-c", + "copy", + str(output_path), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + await process.wait() + return output_path.exists() + except Exception as e: + logger.error(f"Failed to concatenate videos: {e}") + return False + finally: + if list_file.exists(): + list_file.unlink() + + +_drawtext_supported: bool | None = None + + +def is_ffmpeg_drawtext_supported() -> bool: + """Check if ffmpeg supports the drawtext filter (requires libfreetype/libharfbuzz).""" + global _drawtext_supported + if _drawtext_supported is not None: + return _drawtext_supported + try: + res = subprocess.run( + [get_ffmpeg_path(), "-filters"], + capture_output=True, + text=True, + timeout=3, + ) + _drawtext_supported = " drawtext " in res.stdout + except Exception: + _drawtext_supported = False + return _drawtext_supported + + +async def trim_video( + input_path: Path, + start_time: float, + end_time: float | None, + output_path: Path, + fast_copy: bool = True, +) -> bool: + """Trim a video file using ffmpeg. + + Uses fast stream copy (-c copy) by default with fallback to re-encoding. + """ + try: + cmd = [ + get_ffmpeg_path(), + "-y", + "-ss", + str(start_time), + "-i", + str(input_path), + ] + if end_time is not None: + duration = end_time - start_time + cmd.extend(["-t", str(duration)]) + + if fast_copy: + cmd.extend(["-c", "copy", str(output_path)]) + else: + cmd.extend( + [ + "-vf", + "setpts=PTS-STARTPTS", + "-c:v", + "libx264", + "-pix_fmt", + "yuv420p", + str(output_path), + ] + ) + + process = await asyncio.create_subprocess_exec( + *cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + _, stderr = await process.communicate() + if process.returncode == 0 and output_path.exists(): + return True + + if fast_copy: + logger.warning( + f"ffmpeg fast copy trim failed (code {process.returncode})," + f" falling back to re-encoding: {stderr.decode()[:200]}" + ) + return await trim_video(input_path, start_time, end_time, output_path, fast_copy=False) + + logger.error(f"ffmpeg trim failed (code {process.returncode}): {stderr.decode()}") + return False + except Exception as e: + logger.error(f"Failed to trim video: {e}") + return False + + +def cleanup_video_segments(segments: list[Path], keep_path: Path | None = None) -> None: + """Clean up temporary video segments, optionally keeping one path.""" + for segment in segments: + try: + if segment.exists() and segment != keep_path: + segment.unlink() + if segment.parent.exists() and not any(segment.parent.iterdir()): + segment.parent.rmdir() + except Exception: + pass + + +async def compress_video_for_api( + input_path: Path, + target_size_bytes: int = MAX_VIDEO_SIZE_BYTES, + force_compress: bool = False, + start_offset_seconds: float = 0.0, + slowdown_factor: float = 1.0, +) -> Path: + """Compress a video to fit within API size limits using ffmpeg. + + Uses a two-pass approach: + 1. First check if video is already small enough + 2. If not, compress with reduced resolution and bitrate + + Args: + input_path: Path to the input video file + target_size_bytes: Target maximum file size in bytes + force_compress: If True, always perform compression (e.g., to extract + frames at 15fps) + start_offset_seconds: Offset to add to the burned-in timestamp + slowdown_factor: Factor to slow down the video by (e.g. 5.0 for 5x + slower) to increase API frame sampling rate + + Returns: + Path to the compressed video (may be same as input if no compression + needed) + """ + if not input_path.exists(): + raise FileNotFoundError(f"Video file not found: {input_path}") + + current_size = input_path.stat().st_size + logger.info(f"Video size: {current_size / 1024 / 1024:.2f} MB") + + if current_size <= target_size_bytes and not force_compress and slowdown_factor == 1.0: + logger.info( + "Video already within size limit and force_compress=False, no compression needed" + ) + return input_path + + logger.info( + f"Compressing video to fit within {target_size_bytes / 1024 / 1024:.1f}" + f" MB (slowdown={slowdown_factor})" + ) + + output_path = input_path.parent / f"compressed_{input_path.name}" + + # Use Constant Rate Factor (CRF) for high-fidelity UI text readability. + # CRF dynamically allocates bitrates depending on scene motion. + logger.info("Compressing video using CRF=26 for high-fidelity UI text readability.") + + # Compress with ffmpeg: reduce resolution to 720p max, use CRF + # and check if drawtext is supported + vf_parts = ["setpts=PTS-STARTPTS", "scale='min(720,iw)':'-2'"] + if is_ffmpeg_drawtext_supported(): + vf_parts.append( + "drawtext=text='TS\\:" + f" %{{expr_int_format\\:trunc(t+{start_offset_seconds})\\:d}}" + " s':x=w-tw-30:y=120+th+20:fontcolor=white@0.6:fontsize=44:borderw=4:" + "bordercolor=red@0.6:box=1:boxcolor=yellow@0.4:boxborderw=10:font='Sans" + " Bold'" + ) + else: + logger.warning( + "ffmpeg 'drawtext' filter not supported on this system. Skipping" + " burned-in timestamp overlay." + ) + + if slowdown_factor != 1.0: + vf_parts.append(f"setpts={slowdown_factor}*PTS") + + vf_filter = ",".join(vf_parts) + + compress_cmd = [ + get_ffmpeg_path(), + "-y", + "-i", + str(input_path), + "-vf", + vf_filter, + "-r", + "15", # Set framerate to 15 fps + "-c:v", + "libx264", + "-preset", + "fast", + "-crf", + "26", # High quality for UI elements + "-pix_fmt", + "yuv420p", + "-c:a", + "aac", + "-b:a", + "64k", + str(output_path), + ] + + try: + proc = await asyncio.create_subprocess_exec( + *compress_cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + _, stderr = await proc.communicate() + + if proc.returncode != 0: + err_msg = stderr.decode().strip() + logger.error(f"ffmpeg compression failed (code {proc.returncode}): {err_msg}") + if force_compress or slowdown_factor != 1.0: + raise RuntimeError( + f"Video compression/slowdown failed (code {proc.returncode}): {err_msg}" + ) + return input_path # Return original if compression fails and was optional + + new_size = output_path.stat().st_size + logger.info( + f"Compressed: {current_size / 1024 / 1024:.2f} MB -> {new_size / 1024 / 1024:.2f} MB" + ) + + return output_path + + except Exception as e: + logger.error(f"Video compression failed: {e}") + if force_compress or slowdown_factor != 1.0: + raise RuntimeError(f"Video compression/slowdown failed: {e}") + return input_path # Return original if compression fails and was optional + + +async def extract_audio_from_video(input_path: Path) -> Path: + """Extract audio from a video file using ffmpeg. + + Saves as .mp3. + """ + if not input_path.exists(): + raise FileNotFoundError(f"Video file not found: {input_path}") + + output_path = input_path.parent / f"audio_{input_path.stem}.mp3" + logger.info(f"Extracting audio from {input_path} to {output_path}") + + cmd = [ + get_ffmpeg_path(), + "-y", + "-i", + str(input_path), + "-vn", + "-c:a", + "libmp3lame", + "-b:a", + "128k", + str(output_path), + ] + + try: + proc = await asyncio.create_subprocess_exec( + *cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + _, stderr = await proc.communicate() + + if proc.returncode != 0: + err_msg = stderr.decode().strip() + logger.error(f"ffmpeg audio extraction failed: {err_msg}") + raise RuntimeError( + f"ffmpeg audio extraction failed (code {proc.returncode}): {err_msg}" + ) + + return output_path + + except Exception as e: + logger.error(f"Audio extraction failed: {e}") + raise e + + +def extract_keyframes_from_video( + video_path: Path | str, + fps: float = 1.0, + max_frames: int = 45, + max_dimension: int = 1080, +) -> list[tuple[float, bytes]]: + """Extracts keyframes sampled at `fps` intervals. + + Returns: + List of tuples: (timestamp_seconds, jpeg_bytes). + """ + + cap = cv2.VideoCapture(str(video_path)) + if not cap.isOpened(): + logger.error(f"Failed to open video file for keyframe extraction: {video_path}") + return [] + + video_fps = cap.get(cv2.CAP_PROP_FPS) or 30.0 + step = max(1, int(round(video_fps / fps))) + + frames: list[tuple[float, bytes]] = [] + frame_idx = 0 + + while cap.isOpened() and len(frames) < max_frames: + ret, frame = cap.read() + if not ret: + break + + if frame_idx % step == 0: + timestamp_sec = frame_idx / video_fps + h, w = frame.shape[:2] + if max(h, w) > max_dimension: + scale = max_dimension / max(h, w) + frame = cv2.resize( + frame, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA + ) + + _, buf = cv2.imencode(".jpg", frame, [int(cv2.IMWRITE_JPEG_QUALITY), 80]) + frames.append((timestamp_sec, buf.tobytes())) + + frame_idx += 1 + + cap.release() + return frames diff --git a/artemis/utils/visualization.py b/artemis/utils/visualization.py new file mode 100644 index 0000000..e317beb --- /dev/null +++ b/artemis/utils/visualization.py @@ -0,0 +1,693 @@ +# 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 io +import os +import re +from typing import Any + +from artemis.utils.coordinates import parse_swipe_parameters +from PIL import Image, ImageDraw, ImageFont + + +def draw_bounding_boxes( + image_path: str, boxes: list[list[int]], labels: list[str], output_path: str +): + """Draw bounding boxes and labels on an image. + + Args: + image_path: Path to the original image. + boxes: List of boxes, each box is [left, top, right, bottom]. + labels: List of labels corresponding to each box. + output_path: Path to save the annotated image. + """ + if not os.path.exists(image_path): + raise FileNotFoundError(f"Image not found at {image_path}") + + img = Image.open(image_path) + if img.mode in ("RGBA", "P"): + img = img.convert("RGB") + draw = ImageDraw.Draw(img) + + # Use a default font or try to load a system font + try: + font = ImageFont.load_default(size=45) + except Exception: + font = None + + for box, label in zip(boxes, labels): + left, top, right, bottom = box + + # Draw rectangle + draw.rectangle([left, top, right, bottom], outline="red", width=3) + + # Draw label with outline stroke for high contrast + if font: + draw.text( + (left, top), + label, + fill="red", + font=font, + stroke_width=5, + stroke_fill="white", + ) + else: + draw.text( + (left, top), + label, + fill="red", + stroke_width=5, + stroke_fill="white", + ) + img.save(output_path) + + +def draw_dots( + image_path: str, + points: list[list[int]], + labels: list[str], + output_path: str, + radius: int = 10, + color: str = "red", +): + """Draw dots and labels on an image. + + Args: + image_path: Path to the original image. + points: List of points, each point is [x, y]. + labels: List of labels corresponding to each point. + output_path: Path to save the annotated image. + radius: Radius of the dot. + color: Color of the indicator dot and text label. + """ + if not os.path.exists(image_path): + raise FileNotFoundError(f"Image not found at {image_path}") + + img = Image.open(image_path) + if img.mode in ("RGBA", "P"): + img = img.convert("RGB") + draw = ImageDraw.Draw(img) + + try: + font = ImageFont.load_default(size=45) + except Exception: + font = None + + # 1. Draw all background and foreground dots first + for point in points: + x, y = point + + # Draw larger white background dot (2x radius) + white_radius = 2 * radius + draw.ellipse( + [ + x - white_radius, + y - white_radius, + x + white_radius, + y + white_radius, + ], + fill="white", + outline="white", + ) + + # Draw colored indicator dot on top + draw.ellipse( + [x - radius, y - radius, x + radius, y + radius], + fill=color, + outline=color, + ) + + # 2. Draw all text labels on top + for point, label in zip(points, labels): + x, y = point + white_radius = 2 * radius + + # Draw label shifted to prevent overlap with white dot and outline with a white stroke + if font: + draw.text( + (x + white_radius + 4, y - 22), + label, + fill=color, + font=font, + stroke_width=5, + stroke_fill="white", + ) + else: + draw.text( + (x + white_radius + 4, y - 22), + label, + fill=color, + stroke_width=5, + stroke_fill="white", + ) + + img.save(output_path) + + +def parse_bounds(bounds_str: str | None) -> tuple[int, int, int, int] | None: + """Parses standard Android '[left,top][right,bottom]' bounds string.""" + if not bounds_str: + return None + match = re.match(r"\[(\-?\d+),(\-?\d+)\]\[(\-?\d+),(\-?\d+)\]", bounds_str) + if match: + return tuple(map(int, match.groups())) + return None + + +def get_center_coordinates(left: int, top: int, right: int, bottom: int) -> tuple[int, int]: + """Returns pixel center coordinates of a bounding box.""" + return (left + right) // 2, (top + bottom) // 2 + + +def _inject_mutual_occlusion_warnings( + lines: list[str], items_bounds: list[tuple[int, tuple[int, int, int, int]]] +) -> list[str]: + """Detect pairs of elements with >= 50% overlap and append a soft, restrained warning + + to the generated element description lines. + Excludes pure parent-child enclosing bounds to prevent noise on standard + nested containers. + """ + if len(items_bounds) < 2: + return lines + + warnings_map: dict[int, list[int]] = {} + n = len(items_bounds) + + for i in range(n): + idx_i, b_i = items_bounds[i] + l_i, t_i, r_i, bot_i = b_i + area_i = max(1, (r_i - l_i) * (bot_i - t_i)) + + for j in range(i + 1, n): + idx_j, b_j = items_bounds[j] + l_j, t_j, r_j, bot_j = b_j + area_j = max(1, (r_j - l_j) * (bot_j - t_j)) + + # Calculate intersection + inter_l = max(l_i, l_j) + inter_t = max(t_i, t_j) + inter_r = min(r_i, r_j) + inter_bot = min(bot_i, bot_j) + + if inter_l >= inter_r or inter_t >= inter_bot: + continue + + inter_area = (inter_r - inter_l) * (inter_bot - inter_t) + ratio_i = inter_area / area_i + ratio_j = inter_area / area_j + + # To be considered typical parent-child nesting (and NOT occlusion), one node must be inside another, + # AND their centers must be closely aligned (i.e. concentric/enclosing layout). + # If one node overlaps/sits inside another but their centers are significantly offset (e.g. a right-aligned Floating Action Button covering the right end of a left-aligned list item), it is mutual occlusion! + center_i_x = (l_i + r_i) / 2.0 + center_i_y = (t_i + bot_i) / 2.0 + center_j_x = (l_j + r_j) / 2.0 + center_j_y = (t_j + bot_j) / 2.0 + center_dist = ((center_i_x - center_j_x) ** 2 + (center_i_y - center_j_y) ** 2) ** 0.5 + + max_dim = max(r_i - l_i, r_j - l_j, bot_i - t_i, bot_j - t_j) + is_concentric = center_dist < (max_dim * 0.20) + + is_i_inside_j = ( + (l_i >= l_j and t_i >= t_j and r_i <= r_j and bot_i <= bot_j) + and (area_j > area_i * 2.0) + and is_concentric + ) + is_j_inside_i = ( + (l_j >= l_i and t_j >= t_i and r_j <= r_i and bot_j <= bot_i) + and (area_i > area_j * 2.0) + and is_concentric + ) + if is_i_inside_j or is_j_inside_i: + continue + + if ratio_i >= 0.50 or ratio_j >= 0.50: + warnings_map.setdefault(idx_i, []).append(idx_j) + warnings_map.setdefault(idx_j, []).append(idx_i) + + if not warnings_map: + return lines + + new_lines = [] + for line in lines: + match = re.match(r"^\[(\d+)\]", line) + if match: + el_idx = int(match.group(1)) + if el_idx in warnings_map: + overlaps = sorted(list(set(warnings_map[el_idx]))) + overlaps_str = " and ".join(f"[{o}]" for o in overlaps[:2]) + if len(overlaps) > 2: + overlaps_str += f" (+{len(overlaps) - 2} more)" + line = f"{line} (WARNING: may overlap with {overlaps_str}, possible occlusion)" + new_lines.append(line) + return new_lines + + +def format_minimal_list_with_points( + fused_xml: list[dict[str, Any]], width: int = 1080, height: int = 2400 +) -> tuple[str, list[list[int]], list[str]]: + """Formats fused XML hierarchy into a minimal list with sequence numbers, + + and returns physical coordinate points & labels for drawing. + """ + lines = [] + points = [] + labels = [] + idx = 1 + added_items = [] # list of tuples: (text, cx, cy) + items_bounds = [] # list of tuples: (idx, (left, top, right, bottom)) + + def normalize_bounds(left: int, top: int, right: int, bottom: int) -> str: + n_left = int(max(0, min(1000, left * 1000 / width))) + n_top = int(max(0, min(1000, top * 1000 / height))) + n_right = int(max(0, min(1000, right * 1000 / width))) + n_bottom = int(max(0, min(1000, bottom * 1000 / height))) + return f"[{n_left},{n_top}][{n_right},{n_bottom}]" + + def is_duplicate(text_val: str, cx: int, cy: int) -> bool: + text_clean = text_val.strip() + for ex_text, ex_cx, ex_cy in added_items: + if ex_text == text_clean: + dist = ((cx - ex_cx) ** 2 + (cy - ex_cy) ** 2) ** 0.5 + if dist < 8.0: + return True + return False + + for node in fused_xml: + text = node.get("text") or node.get("content-desc") or "" + bounds_str = node.get("bounds") + bounds = parse_bounds(bounds_str) + + ocr_elements = node.get("ocr_elements") + registered = False + + if ocr_elements: + for ocr in ocr_elements: + ocr_bounds = parse_bounds(ocr.get("bounds")) + if ocr_bounds: + left, top, right, bottom = ocr_bounds + cx, cy = get_center_coordinates(left, top, right, bottom) + if is_duplicate(ocr["text"], cx, cy): + continue + + norm_bounds = normalize_bounds(left, top, right, bottom) + lines.append(f"[{idx}] OCR Text: '{ocr['text']}' | Bounds: {norm_bounds}") + items_bounds.append((idx, (left, top, right, bottom))) + + points.append([cx, cy]) + labels.append(str(idx)) + added_items.append((ocr["text"].strip(), cx, cy)) + idx += 1 + registered = True + + if not registered and text.strip() and bounds: + left, top, right, bottom = bounds + cx, cy = get_center_coordinates(left, top, right, bottom) + if is_duplicate(text, cx, cy): + continue + + norm_bounds = normalize_bounds(left, top, right, bottom) + lines.append(f"[{idx}] Text: '{text.strip()}' | Bounds: {norm_bounds}") + items_bounds.append((idx, (left, top, right, bottom))) + + points.append([cx, cy]) + labels.append(str(idx)) + added_items.append((text.strip(), cx, cy)) + idx += 1 + + lines = _inject_mutual_occlusion_warnings(lines, items_bounds) + return "\n".join(lines), points, labels + + +def format_minimal_list_with_elements( + fused_xml: list[dict[str, Any]], width: int = 1080, height: int = 2400 +) -> tuple[str, list[dict[str, Any]], list[str]]: + """Formats fused XML hierarchy into a minimal list with sequence numbers, + + and returns structured element metadata & labels for drawing/validation. + """ + lines = [] + elements = [] + labels = [] + idx = 1 + added_items = [] # list of tuples: (text, cx, cy) + items_bounds = [] # list of tuples: (idx, (left, top, right, bottom)) + + def normalize_bounds(left: int, top: int, right: int, bottom: int) -> str: + n_left = int(max(0, min(1000, left * 1000 / width))) + n_top = int(max(0, min(1000, top * 1000 / height))) + n_right = int(max(0, min(1000, right * 1000 / width))) + n_bottom = int(max(0, min(1000, bottom * 1000 / height))) + return f"[{n_left},{n_top}][{n_right},{n_bottom}]" + + def is_duplicate(text_val: str, cx: int, cy: int) -> bool: + text_clean = text_val.strip() + for ex_text, ex_cx, ex_cy in added_items: + if ex_text == text_clean: + dist = ((cx - ex_cx) ** 2 + (cy - ex_cy) ** 2) ** 0.5 + if dist < 8.0: + return True + return False + + for node in fused_xml: + text = node.get("text") or node.get("content-desc") or "" + bounds_str = node.get("bounds") + bounds = parse_bounds(bounds_str) + + ocr_elements = node.get("ocr_elements") + registered = False + + if ocr_elements: + for ocr in ocr_elements: + ocr_bounds = parse_bounds(ocr.get("bounds")) + if ocr_bounds: + left, top, right, bottom = ocr_bounds + cx, cy = get_center_coordinates(left, top, right, bottom) + if is_duplicate(ocr["text"], cx, cy): + continue + + norm_bounds = normalize_bounds(left, top, right, bottom) + lines.append(f"[{idx}] OCR Text: '{ocr['text']}' | Bounds: {norm_bounds}") + items_bounds.append((idx, (left, top, right, bottom))) + + elements.append( + { + "index": idx, + "center": [cx, cy], + "text": ocr["text"].strip(), + "bounds": [left, top, right, bottom], + "class": node.get("class"), + "resource_id": node.get("resource-id"), + "is_ocr": True, + } + ) + labels.append(str(idx)) + added_items.append((ocr["text"].strip(), cx, cy)) + idx += 1 + registered = True + + if not registered and text.strip() and bounds: + left, top, right, bottom = bounds + cx, cy = get_center_coordinates(left, top, right, bottom) + if is_duplicate(text, cx, cy): + continue + + norm_bounds = normalize_bounds(left, top, right, bottom) + lines.append(f"[{idx}] Text: '{text.strip()}' | Bounds: {norm_bounds}") + items_bounds.append((idx, (left, top, right, bottom))) + + elements.append( + { + "index": idx, + "center": [cx, cy], + "text": text.strip(), + "bounds": [left, top, right, bottom], + "class": node.get("class"), + "resource_id": node.get("resource-id"), + "is_ocr": False, + } + ) + labels.append(str(idx)) + added_items.append((text.strip(), cx, cy)) + idx += 1 + + lines = _inject_mutual_occlusion_warnings(lines, items_bounds) + return "\n".join(lines), elements, labels + + +def crop_and_annotate_target( + image_bytes: bytes, + coords: list[int], + crop_size: int | None = 256, + dot_radius: int = 8, +) -> bytes: + """Crops a square region centered around coords, and draws a red dot with black outline at coords. + + If crop_size is None, returns the full annotated image. + """ + img = Image.open(io.BytesIO(image_bytes)) + width, height = img.size + x, y = coords + + # Ensure coordinates are within image boundaries + x = max(0, min(x, width - 1)) + y = max(0, min(y, height - 1)) + + if crop_size is not None: + # Calculate crop boundaries (centered around x, y; clamped to image boundaries) + w_crop = min(crop_size, width) + h_crop = min(crop_size, height) + + x_min = max(0, min(x - w_crop // 2, width - w_crop)) + y_min = max(0, min(y - h_crop // 2, height - h_crop)) + x_max = x_min + w_crop + y_max = y_min + h_crop + + # Crop + cropped_img = img.crop((x_min, y_min, x_max, y_max)) + + # Local coordinates within the cropped frame + local_x = x - x_min + local_y = y - y_min + else: + cropped_img = img + local_x = x + local_y = y + + # Draw solid red circle with a black outline + draw = ImageDraw.Draw(cropped_img) + draw.ellipse( + [ + local_x - dot_radius, + local_y - dot_radius, + local_x + dot_radius, + local_y + dot_radius, + ], + fill="red", + outline="black", + width=2 if crop_size is None else 1, + ) + + out_buf = io.BytesIO() + if crop_size is None: + cropped_img.save(out_buf, format="JPEG", quality=85) + else: + cropped_img.save(out_buf, format="PNG") + return out_buf.getvalue() + + +def _resolve_coordinates(coord: Any, width: int, height: int) -> tuple[int, int] | None: + """Helper to safely parse coordinates into pixel space.""" + if not isinstance(coord, (list, tuple)) or len(coord) < 2: + return None + try: + x, y = float(coord[0]), float(coord[1]) + # If coordinates are normalized (0.0 <= x <= 1.0) + if 0.0 <= x <= 1.0 and 0.0 <= y <= 1.0: + px = int(round(x * width)) + py = int(round(y * height)) + elif 0 <= x <= 1000 and 0 <= y <= 1000 and (width > 1000 or height > 1000): + # Android normalized 1000 coordinate space + px = int(round(x * width / 1000.0)) + py = int(round(y * height / 1000.0)) + else: + px = int(round(x)) + py = int(round(y)) + px = max(0, min(width - 1, px)) + py = max(0, min(height - 1, py)) + return px, py + except Exception: + return None + + +def draw_action_overlay_on_image( + image_bytes: bytes, + action_name: str, + action_args: dict[str, Any], +) -> bytes: + """Draws visual action indicators (tap ripples, sequence numbers, swipe arrows, focus boxes) + + on top of the before-action screenshot. + + Args: + image_bytes: Raw JPEG/PNG image bytes of the screenshot. + action_name: Action name (e.g. 'click', 'click_sequence', 'swipe', 'input_text'). + action_args: Action parameters dictionary. + + Returns: + Annotated JPEG image bytes, or original image_bytes on error. + """ + import math + + if not image_bytes: + return image_bytes + + try: + img = Image.open(io.BytesIO(image_bytes)) + if img.mode in ("RGBA", "P"): + img = img.convert("RGB") + width, height = img.size + draw = ImageDraw.Draw(img) + + # Scale indicator sizes dynamically based on screen resolution + base_dim = min(width, height) + dot_radius = max(18, int(base_dim * 0.025)) + line_width = max(4, int(base_dim * 0.006)) + + action = (action_name or "").lower() + + # 1. Single Click / Tap + if action in ("click", "tap", "click_coordinate", "long_click", "press"): + coord = ( + action_args.get("target") + or action_args.get("coordinates") + or action_args.get("point") + ) + pt = _resolve_coordinates(coord, width, height) + if pt: + x, y = pt + # Outer white aura + draw.ellipse( + [ + x - dot_radius * 1.5, + y - dot_radius * 1.5, + x + dot_radius * 1.5, + y + dot_radius * 1.5, + ], + outline="white", + width=line_width + 2, + ) + # Red ripple circle + draw.ellipse( + [x - dot_radius, y - dot_radius, x + dot_radius, y + dot_radius], + fill="red", + outline="white", + width=line_width, + ) + # Crosshair + cross_len = int(dot_radius * 1.8) + draw.line([(x - cross_len, y), (x + cross_len, y)], fill="white", width=line_width) + draw.line([(x, y - cross_len), (x, y + cross_len)], fill="white", width=line_width) + + # 2. Click Sequence + elif action in ("click_sequence", "tap_sequence"): + sequence = action_args.get("sequence") or action_args.get("targets") or [] + if isinstance(sequence, list): + for idx, coord in enumerate(sequence, 1): + pt = _resolve_coordinates(coord, width, height) + if pt: + x, y = pt + draw.ellipse( + [x - dot_radius, y - dot_radius, x + dot_radius, y + dot_radius], + fill="red", + outline="white", + width=line_width, + ) + label = str(idx) + try: + font = ImageFont.load_default(size=int(dot_radius * 1.2)) + draw.text( + (x - dot_radius // 3, y - dot_radius // 2), + label, + fill="white", + font=font, + stroke_width=2, + stroke_fill="black", + ) + except Exception: + draw.text( + (x - 4, y - 6), + label, + fill="white", + ) + + # 3. Swipe / Drag / Scroll + elif action in ("swipe", "drag", "scroll"): + kind, target, _ = parse_swipe_parameters(action_args) + start_pt, end_pt = None, None + + if kind == "coords" and isinstance(target, list) and len(target) == 4: + start_pt = _resolve_coordinates(target[:2], width, height) + end_pt = _resolve_coordinates(target[2:], width, height) + elif kind == "direction" and isinstance(target, str): + direction = target.lower() + cx, cy = width // 2, height // 2 + dist_y = int(height * 0.35) + dist_x = int(width * 0.35) + if "up" in direction: + start_pt = (cx, cy + dist_y // 2) + end_pt = (cx, cy - dist_y // 2) + elif "down" in direction: + start_pt = (cx, cy - dist_y // 2) + end_pt = (cx, cy + dist_y // 2) + elif "left" in direction: + start_pt = (cx + dist_x // 2, cy) + end_pt = (cx - dist_x // 2, cy) + elif "right" in direction: + start_pt = (cx - dist_x // 2, cy) + end_pt = (cx + dist_x // 2, cy) + + if start_pt and end_pt: + sx, sy = start_pt + ex, ey = end_pt + + draw.ellipse( + [ + sx - dot_radius // 2, + sy - dot_radius // 2, + sx + dot_radius // 2, + sy + dot_radius // 2, + ], + fill="yellow", + outline="black", + width=2, + ) + draw.line([(sx, sy), (ex, ey)], fill="white", width=line_width + 4) + draw.line([(sx, sy), (ex, ey)], fill="red", width=line_width) + + dx = ex - sx + dy = ey - sy + angle = math.atan2(dy, dx) + arrow_len = max(24, int(base_dim * 0.04)) + arrow_angle = math.pi / 6 + + x1 = ex - arrow_len * math.cos(angle - arrow_angle) + y1 = ey - arrow_len * math.sin(angle - arrow_angle) + x2 = ex - arrow_len * math.cos(angle + arrow_angle) + y2 = ey - arrow_len * math.sin(angle + arrow_angle) + + draw.polygon([(ex, ey), (x1, y1), (x2, y2)], fill="red", outline="white") + + # 4. Input Text + elif action in ("input_text", "type"): + coord = action_args.get("target") or action_args.get("coordinates") + if coord: + pt = _resolve_coordinates(coord, width, height) + if pt: + x, y = pt + box_w = max(60, int(base_dim * 0.15)) + box_h = max(30, int(base_dim * 0.06)) + draw.rectangle( + [x - box_w, y - box_h, x + box_w, y + box_h], + outline="cyan", + width=line_width, + ) + + out_buf = io.BytesIO() + img.save(out_buf, format="JPEG", quality=85) + return out_buf.getvalue() + + except Exception: + return image_bytes diff --git a/config/artemis.jsonc b/config/artemis.jsonc new file mode 100644 index 0000000..7a73e41 --- /dev/null +++ b/config/artemis.jsonc @@ -0,0 +1,217 @@ +{ + // ============================================================================== + // โ˜• ARTEMIS Autonomous Agent & Framework Configuration + // ============================================================================== + + // ------------------------------------------------------------------------------ + // 1. Global Default Model (Applied to all sub-agents unless overridden below) + // ------------------------------------------------------------------------------ + "default": { + "provider": "google", + "model": "gemini-3.6-flash", + "thinking_level": "medium", + "fallback": { + "provider": "google", + "model": "gemini-3.5-flash", + "thinking_level": "medium" + } + }, + + // ------------------------------------------------------------------------------ + // 2. Ready-to-Use Presets (Switch easily in code or CLI) + // ------------------------------------------------------------------------------ + "presets": { + "gemini-flagship": { + "provider": "google", + "model": "gemini-3.6-flash", + "fallback": { "provider": "google", "model": "gemini-3.5-flash" } + }, + "gemini-flash": { + "provider": "google", + "model": "gemini-2.5-flash", + "fallback": { "provider": "google", "model": "gemini-2.0-flash" } + }, + "openai-gpt4o": { + "provider": "openai", + "model": "gpt-4o", + "fallback": { "provider": "openai", "model": "gpt-4o-mini" } + }, + "cost-saving": { + "provider": "google", + "model": "gemini-3.5-flash-lite", + "fallback": { "provider": "google", "model": "gemini-3.1-flash-lite" } + }, + "local-ollama": { + "provider": "openai", + "model": "llama3.2-vision", + "fallback": { "provider": "openai", "model": "llama3.2-vision" } + } + }, + + // ------------------------------------------------------------------------------ + // 3. Agent & Sub-Agent Node Customizations (Optional - Inherits from 'default') + // ------------------------------------------------------------------------------ + "nodes": { + // ๐Ÿง  Planner: High-level goal decomposition and strategic step planning + "planner": { + "thinking_level": "medium" + }, + + // ๐ŸŽฏ Operator: Screen perception, element interaction, and direct execution + "operator": { + "thinking_level": "medium", + "include_thoughts": true + }, + + // ๐ŸŽฌ Video Analyzer: Multimodal screen recording verification and temporal analysis + // NOTE: Dual-path engine automatically routes: + // - Native Gemini Files API for Gemini models with Google API keys (long video file stream) + // - Universal Multi-Model Engine for Claude 3.5/3.7, GPT-4o, OpenRouter, Vertex AI (dynamic keyframe extraction + multimodal tool calling) + "video_analyzer": {}, + + // ๐Ÿ‘๏ธ Object Detector & Explorer (Flash Mode): Visual coordinates grounding and bounding boxes + // CRITICAL REQUIREMENT: Standard LLMs (including standard Gemini Flash, GPT-4o, Claude 3.7) lack sub-pixel spatial coordinate fine-tuning. + // To achieve accurate [x, y] point and bounding box localization, this node MUST use specialized Gemini ER (Embodied Reasoning / Robotics) models: + // e.g., "gemini-robotics-er-2-preview". Non-ER models will fail spatial coordinate detection. + "object_detector": { + "model": "gemini-robotics-er-2-preview", + "fallback": { + "provider": "google", + "model": "gemini-robotics-er-2-preview" + } + }, + + // ๐Ÿ” Explorer: Deep element grounding and multi-candidate visual search + // NOTE: + // - "flash" version delegates directly to object_detector (requires Gemini ER model above). + // - "pro" / "ultra" versions run multi-turn ReAct loops (supports both Native Gemini and Universal Multi-Model LangChain engines). + "explorer": { + "thinking_level": "medium" + }, + + // ๐Ÿ›ก๏ธ Validator & Failure Analyzer: Self-healing, action verification & recovery + "validator_failure_analyzer": { + "thinking_level": "medium" + }, + + // ๐Ÿ“‹ Summarizer: Aggregates step history and generates structured outcome + "summarizer": {}, + + // ๐Ÿ”ฝ Hopper: Fast lightweight package name locator + "hopper": { + "model": "gemini-3.5-flash-lite", + "fallback": { + "provider": "google", + "model": "gemini-3.1-flash-lite" + } + }, + + // ๐Ÿ“ Outputter: Synthesizes execution trace into final output report + "outputter": {}, + + // ๐Ÿชต Log Analyzer: Error log diagnosis with search grounding + "log_analyzer": { + "enable_grounding": true + }, + + // ๐Ÿฉบ Diagnoser Expert: Failure root cause diagnosis with search grounding + "diagnoser_expert": { + "enable_grounding": true + }, + // ๐Ÿ›๏ธ Committee Council Members + "planner_avatar": {}, + "history_analyzer_expert": {}, + // ๐Ÿ”Ž Checker: Visual subgoal verification and rollback + "checker": {} + }, + + // ------------------------------------------------------------------------------ + // 4. Runtime Mode & Agent Capabilities Configuration (Dual-Engine: Flash vs. Pro) + // ------------------------------------------------------------------------------ + "agent": { + // ========================================================================== + // โšก 4.1 FLASH Profile Options (Reactive Single-Agent Lightweight Loop) + // ========================================================================== + // Best for: Simple, fast, deterministic tasks (< 30 UI steps). + // Execution: Observe-Think-Act single LLM reactive loop without graph overhead. + "flash": { + "max_turns": 30, // Maximum reactive turns before stopping + "explorer_mode": "flash", // Perception mode for Flash runner ('flash' 1-shot detection) + // ๐Ÿง  Step Summarizer: Asynchronous visual context compressor + // Summarizes intermediate steps in background without blocking the runner, + // replacing pruned images with strictly objective, high-density state summaries. + "step_summarizer": { + "enabled": true, // Whether background step summarization is active + "model": "gemini-3.5-flash-lite", // Lightweight model for fast background inference + "prune_history_xml": true // Prune outdated heavy XML trees from historical steps + } + }, + // ========================================================================== + // ๐Ÿš€ 4.2 PRO Profile Options (Deep Multi-Agent Graph Closed-Loop) + // ========================================================================== + // Best for: Complex, exploratory, dynamic multi-branch tasks. + // Execution: LangGraph multi-node graph with autonomous planning, validation, and self-healing. + "pro": { + // ๐Ÿงญ Explorer: Perception sub-agent under Pro profile ('flash', 'pro', or 'ultra') + "explorer": { + "mode": "flash", // Default mode for Pro nodes (Operator/Validator) + "caching": true // Context caching for multi-turn Pro/Ultra visual reasoning + }, + // ๐ŸŽฏ Planner Validation: Async verification when task plan milestones change + "planner_validation": { + "enabled": false, + "similarity_threshold": 0.85 + }, + // ๐Ÿ›๏ธ Committee: Multi-Agent Council debate tool for Operator + "committee": { + "enabled": false, + "debate_rounds": 2 + }, + // ๐Ÿ”Ž Checker: Visual subgoal verification and automatic checkpoint rollback + "checker": { + "enabled": false, + "max_iterations": 20, + "max_chat_rounds": 4 + }, + // ๐ŸŽฌ Video Analyzer & Screen Recording: Continuous screen recording via scrcpy and temporal analysis + "video_analyzer": { + "enabled": true, // Whether continuous video recording and video tools are active (default: auto-detected) + "enable_ledger": true + } + }, + // ========================================================================== + // ๐ŸŒ 4.3 Shared Capabilities & Global Overrides + // ========================================================================== + // ๐ŸŽฌ Screen Recording & Video Analyzer Configuration + "video_analyzer": { + "enabled": true, // Whether continuous screen recording via scrcpy is active (default: auto-detected) + "enable_ledger": true + }, + // ๐Ÿ“ Outputter: Post-execution report synthesis and structured schema extraction + // Reviews action history, notes, and verifies goal achievement to synthesize final output. + "outputter": { + "enabled": true, // Whether Outputter is mounted to synthesize final outputs + "force_synthesis": false // Force Outputter synthesis even without explicit output schema + }, + // ๐Ÿงญ Global / Per-Agent Explorer Version Overrides + "explorer": { + "default_version": "flash", // Fallback version when not specified + "flash_mode": "flash", // Version for Flash runner + "pro_mode": "flash", // Version for Pro graph runner + "caching": true + }, + "explorer_versions": { + "operator": "flash", + "validator": "flash" + }, +// ๐Ÿšซ Per-agent Tool Blacklists + "blacklisted_tools": { + "explorer": [ + "ask_image_processor", + "get_ocr_list", + "inspect_region" + ] + } + } +} + diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..e4f11d9 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,218 @@ +# 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. + +[project] +name = "artemis" +version = "3.6.3" +description = "Artemis Android Agent" + +authors = [ + { name = "somew1nd", email = "wfq559@outlook.com" }, +] + +requires-python = ">=3.12" + +dependencies = [ + # Orchestration & Framework + "langgraph>=1.0.2,<2.0.0", + "langchain>=1.0.0", + "langchain-core>=1.0.0", + "langchain-community>=0.4.1", + "langchain-mcp-adapters>=0.2.0", + "mcp>=1.26.0", + # LLM Providers & GenAI SDKs + "google-genai>=0.1.1", + "langchain-google-genai>=4.0.0", + "langchain-google-vertexai>=3.0.0", + "langchain-openai>=1.0.0", + # Configuration, Models & Validation + "pydantic>=2.7.0", + "pydantic-settings>=2.10.1", + "python-dotenv>=1.0.1", + "jinja2>=3.1.5", + "typing-extensions>=4.12.0", + # Android Automation & Device Subsystem + "adbutils>=2.12.0", + "uiautomator2>=3.5.0", + # CLI & Rich Terminal UI + "typer>=0.12.0", + "rich>=13.7.0", + "inquirer>=3.4.0", + "colorama>=0.4.6", + # Web, API & Networking + "fastapi>=0.111.0", + "uvicorn[standard]>=0.30.1", + "httpx>=0.28.1", + "requests>=2.31.0", + "websockets>=12.0", + "sseclient-py>=1.8.0", + # Image, Video & Multimodal Perception + "pillow>=10.0.0", + "opencv-python>=4.9.0.80", + "imageio-ffmpeg>=0.5.1", + "matplotlib>=3.8.0", + "numpy>=1.26.0", + "scipy>=1.13.0", + # System, Utilities & Telemetry + "psutil>=5.9.0", + "retry>=0.9.2", + "posthog>=3.5.0", + "uuid-utils>=0.9.0", + "jupyter-client>=8.6.0", + "ipykernel>=6.29.0", + "langchain-anthropic>=1.5.4", +] + +[project.optional-dependencies] +dev = [ + "ruff>=0.15.2", + "pytest>=8.4.1", + "pytest-asyncio>=0.23.0", + "pytest-cov>=5.0.0", + "syrupy>=4.0.0", + "pyright>=1.1.389", + "pre-commit>=3.5.0", +] + +[project.scripts] +artemis = "artemis.interfaces.cli.main:cli" +artemis-admin = "apps.admin_console.server:main" + + +[build-system] +requires = ["setuptools>=61.0", "Cython"] +build-backend = "setuptools.build_meta" + +[tool.ruff] +# Exclude a variety of commonly ignored directories. +exclude = [ + ".bzr", + ".direnv", + ".eggs", + ".git", + ".git-rewrite", + ".hg", + ".ipynb_checkpoints", + ".mypy_cache", + ".nox", + ".pants.d", + ".pyenv", + ".pytest_cache", + ".pytype", + ".ruff_cache", + ".svn", + ".tox", + ".venv", + ".vscode", + "__pypackages__", + "_build", + "buck-out", + "build", + "dist", + "node_modules", + "site-packages", + "venv", +] + +# Same as Black. +line-length = 100 +indent-width = 4 + +target-version = "py312" + +[tool.ruff.lint] +# Enable Pyflakes (`F`) and a subset of the pycodestyle (`E`) codes by default. +# Also enable TID (flake8-tidy-imports) to enforce absolute imports +select = ["E", "F", "TID", "UP"] +ignore = ["E501"] + +# Allow fix for all enabled rules (when `--fix`) is provided. +fixable = ["ALL"] +unfixable = [] + +# Allow unused variables when underscore-prefixed. +dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$" + +[tool.ruff.lint.flake8-tidy-imports] +# Ban relative imports +ban-relative-imports = "all" + +[tool.ruff.format] +# Like Black, use double quotes for strings. +quote-style = "double" + +# Like Black, indent with spaces, rather than tabs. +indent-style = "space" + +# Like Black, respect magic trailing commas. +skip-magic-trailing-comma = false + +# Like Black, automatically detect the appropriate line ending. +line-ending = "auto" + +[tool.pytest.ini_options] +# Add project root and apps to Python path for clean absolute imports +pythonpath = [ + ".", + "apps", + "apps/admin_console", +] + + +# Specify test discovery patterns +testpaths = [ + "tests" +] +python_files = [ + "*.test.py", + "test_*.py", + "*_test.py" +] +python_functions = [ + "test_*" +] +python_classes = [ + "Test*" +] +# Add useful options +addopts = [ + "-v", # verbose output + "--tb=short", # shorter traceback format + "--strict-markers", # strict marker checking + "--disable-warnings", # disable warnings for cleaner output +] + +markers = [ + "android: tests requiring Android device or emulator", +] + +[tool.pylint.master] +ignore = [".venv", "build", "dist"] +recursive = "yes" + +[tool.pylint.messages_control] +disable = [ + "C0114", # missing-module-docstring + "R0801", # duplicate-code +] + +[dependency-groups] +dev = [ + "ruff>=0.15.2", + "pytest>=8.4.1", + "pytest-asyncio>=1.2.0", + "pytest-cov>=5.0.0", + "pyright>=1.1.389", + "pre-commit>=3.5.0", +] diff --git a/pyrightconfig.json b/pyrightconfig.json new file mode 100644 index 0000000..c6f60f5 --- /dev/null +++ b/pyrightconfig.json @@ -0,0 +1,3 @@ +{ + "exclude": [".venv/**"] +} diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..43041dc --- /dev/null +++ b/setup.py @@ -0,0 +1,104 @@ +# 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 os +from pathlib import Path +from setuptools import Extension, find_namespace_packages, setup +from setuptools.command.build_py import build_py + +USE_CYTHON = os.environ.get("USE_CYTHON", "0") == "1" + +ext_modules = [] +cmdclass = {} + +if USE_CYTHON: + from Cython.Build import cythonize + + print(f"USE_CYTHON: {USE_CYTHON}") + + extensions = [] + src_dir = Path("artemis") + + # 1. Recursively find all .py files inside the src/ directory + for py_file in src_dir.rglob("*.py"): + # Get the path relative to src/ (e.g., mypackage/subfolder/module.py) + rel_path = py_file.relative_to(src_dir) + + # Convert path into a Python module name (e.g., artemis.subfolder.module) + mod_parts = ("artemis",) + rel_path.with_suffix("").parts + module_name = ".".join(mod_parts) + + # PREVENT CYTHON CRASH: Skip rogue __init__.py files at the root of src/ + if module_name in ( + "artemis.__init__", + "artemis.main", + "artemis.mcp.adb_server", + "artemis.mcp.xml_search_server", + ): + print(f"Skipping init/main file from Cython: {py_file}") + continue + + # 2. Create an Extension for this specific file + extensions.append(Extension(module_name, [str(py_file)])) + + # Compile all found extensions + ext_modules = cythonize( + extensions, + compiler_directives={"language_level": "3"}, + build_dir="build/cythonized", + ) + + # 3. Prevent packaging original .py source files (except main and __init__) and intermediate .c files + class BuildPyExcludeSource(build_py): + def find_package_modules(self, package, package_dir): + modules = super().find_package_modules(package, package_dir) + return [ + m + for m in modules + if m[1] in ("main", "__init__", "adb_server", "xml_search_server") + ] + + def find_data_files(self, package, package_dir): + data_files = super().find_data_files(package, package_dir) + return [f for f in data_files if not f.endswith(".c")] + + cmdclass["build_py"] = BuildPyExcludeSource + +setup( + name="artemis", + version="3.6.3", + packages=find_namespace_packages( + include=[ + "artemis", + "artemis.*", + "apps", + "apps.*", + "admin_console", + "admin_console.*", + ] + ), + package_dir={ + "": ".", + "admin_console": "apps/admin_console", + }, + package_data={ + "artemis": ["**/*.json", "**/*.md"], + "apps.admin_console": ["index.html", "**/*.json", "**/*.md"], + "admin_console": ["index.html", "**/*.json", "**/*.md"], + }, + include_package_data=True, + ext_modules=ext_modules, + cmdclass=cmdclass, + exclude_package_data={"": ["*.c", "*.pyx", "*.pxd", "*.cpp"]}, +)