chenmingxuan/Artemis-Eapil
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 32 Wiki Activity
Files
codex/v2.0.0-plugins
Artemis-Eapil/apps/admin_console/replay_manager.py
T
chenmingxuan 2ddec53954
CI / Python 3.12 / ubuntu-latest (push) Canceled after 0s
Details
CI / Python 3.12 / windows-latest (push) Canceled after 0s
Details
CI / Frontend tests and production build (push) Canceled after 0s
Details
CI / Build and smoke-test Python packages (push) Canceled after 0s
Details
feat: add Chinese UI, Responses models, and mirrored Windows setup 
Localize the web console and persist custom model and capability settings.
Use OpenAI Responses exclusively and move runtime limits into environment configuration.
Bootstrap uv and the virtual environment on Windows, support optional China mirrors,
and package source distributions without environments, build output, or credentials.

Update regression tests and repository quality checks for the new behavior.
2026-09-22 19:16:17 +08:00

2503 lines
98 KiB
Python
Raw Permalink Blame History

# 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 shutil
import sqlite3
import sys
import traceback
from typing import Any
from uuid import UUID
WORKSPACE_ROOT = Path(__file__).parent.parent.resolve()
# Try to import settings from artemis
try:
from artemis.config import settings
TRACES_PATH = Path(settings.TRACES_PATH)
except ImportError:
# Fallback to default path if run outside project environment or imports fail
TRACES_PATH = WORKSPACE_ROOT / "traces"
DB_PATH = TRACES_PATH / "data_engine.db"
IMAGES_DIR = TRACES_PATH / "images"
try:
from artemis.context import ArtemisContext
from artemis.graph.state import State
except ImportError:
ArtemisContext = Any
State = Any
REPLAY_TOOLS_CONFIG = {
"ask_explorer": {
"display_name": "Ask Explorer",
"description": (
"Visual parsing agent that executes visual ReAct search loops using multi-modal models."
),
"module": "artemis.tools.explorer_tool",
"function": "_run_explorer_logic",
"agent_name": "explorer",
"denylist_args": ["ctx", "state"],
"fallback_mappings": {},
"llm_span_name": "gemini_explorer_call",
},
"ask_image_processor": {
"display_name": "Ask Image Processor",
"description": (
"Visual parsing agent that writes Python code using OpenCV and"
" computer vision to analyze and modify screen images."
),
"module": "artemis.tools.image_processor_tool",
"function": "_run_image_processor_logic",
"agent_name": "image_processor",
"denylist_args": ["ctx", "state"],
"fallback_mappings": {},
"llm_span_name": "gemini_image_processor_call",
},
# Future tools can easily be registered here using plain dictionaries:
# "ask_validator": {
# "display_name": "Ask Validator",
# "description": "Validation agent that checks if a step or action was completed successfully.",
# "module": "artemis.tools.validator_tool",
# ...
# }
}
class ReplayManager:
"""Manages sandboxed execution and diagnostics of visual agent steps.
The ReplayManager is the core orchestrator for the Artemis step replay system.
It reconstructs historical execution states (SQLite database records, trace
histories, and screenshots) and replays a single agent step in an isolated,
deterministic sandbox environment while interacting live with a physical
Android device and OpenAI Responses.
Key Responsibilities:
- Reconstructs and preloads historical SQLite states up to a specific
step.
- Establishes sandboxed `ArtemisContext` and `DataEngine` instances.
- Integrates with physical devices via ADB and UIAutomator.
- Records Responses model calls to capture thoughts and raw outputs.
- Reconciles live and preloaded traces into unified chronological trees.
Attributes:
workspace_root (Path): The absolute path to the root of the project
workspace.
replay_base_dir (Path): Base directory containing replay data and
outputs.
test_data_dir (Path): Directory containing chunked historical step
traces.
test_outputs_dir (Path): Sandbox directory where replay outputs are
written.
db_path (Path): Path to the master Data Engine SQLite database.
traces_path (Path): Path to the master traces directory.
images_dir (Path): Path to the master images directory.
"""
def __init__(
self,
workspace_root: Path = WORKSPACE_ROOT,
device_id: str = None,
original_db_path: Path = None,
init_device: bool = False,
):
"""Initializes the ReplayManager and autodetects the connected Android device.
Args:
workspace_root (Path, optional): The absolute path to the workspace
root. Defaults to WORKSPACE_ROOT.
device_id (str, optional): The serial ID of the target Android
device. If not specified, autodetects the first connected device.
original_db_path (Path, optional): Path to the master SQLite
database. Defaults to DB_PATH.
init_device (bool, optional): Whether to initialize the ADB device
connection on startup.
"""
self.workspace_root = workspace_root
self.traces_path = TRACES_PATH
self.replay_base_dir = self.traces_path / "replay"
self.test_data_dir = self.replay_base_dir / "data"
self.test_outputs_dir = self.replay_base_dir / "outputs"
self.db_path = Path(original_db_path) if original_db_path else DB_PATH
self.images_dir = IMAGES_DIR
# Device connection and autodetection
self.device_id = None
self.w = None
self.h = None
self.adb = None
self.ui_client = None
if init_device:
try:
self._init_device(device_id)
except Exception as e:
print(
"Warning: Connected device could not be initialized during"
f" ReplayManager startup: {e}"
)
def _init_device(self, device_id: str = None):
"""Initializes connection to the target device and queries screen metrics."""
from adbutils import AdbClient
try:
from artemis.clients.ui_automator_client import UIAutomatorClient
except ImportError:
raise ImportError(
"Failed to import UIAutomatorClient. Ensure artemis package is installed in path."
)
self.adb = AdbClient(host="localhost", port=5037)
try:
devices = self.adb.device_list()
except Exception as adb_err:
raise ConnectionError(
f"Failed to query device list from ADB server: {adb_err}"
) from adb_err
if not devices:
raise ConnectionError(
"No active ADB devices connected. Please connect an Android device via ADB."
)
if device_id:
matched_device = next((d.serial for d in devices if d.serial == device_id), None)
if not matched_device:
raise ConnectionError(
f"Requested device '{device_id}' is not connected."
f" Connected devices: {[d.serial for d in devices]}"
)
self.device_id = matched_device
else:
self.device_id = devices[0].serial
self.ui_client = UIAutomatorClient(device_id=self.device_id)
ui_data = self.ui_client.get_screen_data()
self.w, self.h = ui_data.width, ui_data.height
print(f"Connected to device: {self.device_id} ({self.w}x{self.h})")
def chunk_session_traces(self, session_id: str, output_dir: Path = None) -> Path:
"""Chunks traces, screenshots, and metadata for a specific execution session.
Queries the master SQLite database and exports step-by-step
subdirectories
(e.g., `step_00/`, `step_01/`) inside the output directory.
Args:
session_id: The UUID string of the session to extract and chunk.
output_dir: Optional path to the directory where chunked results
should be saved. Defaults to `self.test_data_dir /
f"{session_id}_chunked"`.
"""
import sqlite3
import json
import shutil
db_path = self.db_path
traces_dir = self.traces_path
if not db_path.exists():
raise FileNotFoundError(f"Master database not found at {db_path}")
if not output_dir:
output_dir = self.test_data_dir / f"{session_id}_chunked"
else:
output_dir = Path(output_dir)
if output_dir.exists() or output_dir.is_symlink():
# Safety check: only delete directories that are named *_chunked
if output_dir.name.endswith("_chunked"):
if output_dir.is_symlink():
output_dir.unlink()
elif (
self.test_data_dir in output_dir.parents
or output_dir.parent == self.test_data_dir
or output_dir.parent.name == "inputs"
):
shutil.rmtree(output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
print(f"Chunking traces for session: {session_id}")
print(f"Using database: {db_path}")
print(f"Output directory: {output_dir}")
conn = sqlite3.connect(db_path)
conn.row_factory = sqlite3.Row
cursor = conn.cursor()
cursor.execute("SELECT * FROM sessions WHERE session_id = ?", (session_id,))
session_row = cursor.fetchone()
cursor.execute(
"SELECT * FROM steps WHERE session_id = ? ORDER BY step_number ASC",
(session_id,),
)
steps = cursor.fetchall()
print(f"Found {len(steps)} steps for session {session_id}")
if not session_row and not steps:
conn.close()
raise ValueError(f"Session {session_id} not found in database.")
for step in steps:
step_number = step["step_number"]
step_id = step["step_id"]
step_dir = output_dir / f"step_{step_number:02d}"
step_dir.mkdir(parents=True, exist_ok=True)
# 1. Process and save the step's primary metadata (step.json)
step_data = {}
for key in step.keys():
val = step[key]
if (
key
in (
"action_taken",
"last_execution_result",
"extra_metadata",
)
and val
):
try:
step_data[key] = json.loads(val)
except Exception:
step_data[key] = val
else:
step_data[key] = val
step_json_path = step_dir / "step.json"
with open(step_json_path, "w", encoding="utf-8") as f:
json.dump(step_data, f, indent=2, ensure_ascii=False)
# 2. Retrieve and save trace logs recorded during this step (traces.json)
cursor.execute(
"SELECT * FROM traces WHERE step_id = ? ORDER BY timestamp ASC",
(step_id,),
)
traces = cursor.fetchall()
traces_data = []
for trace_row in traces:
trace_dict = {}
for key in trace_row.keys():
val = trace_row[key]
if key == "payload" and val:
try:
trace_dict[key] = json.loads(val)
except Exception:
trace_dict[key] = val
else:
trace_dict[key] = val
traces_data.append(trace_dict)
traces_json_path = step_dir / "traces.json"
with open(traces_json_path, "w", encoding="utf-8") as f:
json.dump(traces_data, f, indent=2, ensure_ascii=False)
# 3. Copy the step's screenshots (pre-action and post-action screenshots)
pre_image_name = step_data.get("pre_image_name")
post_image_name = step_data.get("post_image_name")
if pre_image_name and post_image_name == pre_image_name:
post_image_name = None
images_dir = traces_dir / "images"
if pre_image_name:
src_pre = images_dir / f"{pre_image_name}.jpg"
if src_pre.exists():
shutil.copy2(src_pre, step_dir / "pre.jpg")
else:
print(f"Warning: Pre-action screenshot not found at {src_pre}")
if post_image_name:
src_post = images_dir / f"{post_image_name}.jpg"
if src_post.exists():
shutil.copy2(src_post, step_dir / "post.jpg")
else:
print(f"Warning: Post-action screenshot not found at {src_post}")
# 4. Save metadata of the pre-action screenshot (pre_image_meta.json)
if pre_image_name:
cursor.execute(
"SELECT * FROM images WHERE image_name = ?",
(pre_image_name,),
)
img_row = cursor.fetchone()
if img_row:
img_meta = {}
for key in img_row.keys():
val = img_row[key]
if key in ("ocr_result", "ui_tree", "extra_metadata") and val:
try:
img_meta[key] = json.loads(val)
except Exception:
img_meta[key] = val
else:
img_meta[key] = val
pre_img_meta_path = step_dir / "pre_image_meta.json"
with open(pre_img_meta_path, "w", encoding="utf-8") as f:
json.dump(img_meta, f, indent=2, ensure_ascii=False)
else:
print(
f"Warning: Image metadata row for {pre_image_name} not"
" found in images table."
)
# 5. Save metadata of the post-action screenshot (post_image_meta.json)
if post_image_name:
cursor.execute(
"SELECT * FROM images WHERE image_name = ?",
(post_image_name,),
)
img_row = cursor.fetchone()
if img_row:
img_meta = {}
for key in img_row.keys():
val = img_row[key]
if key in ("ocr_result", "ui_tree", "extra_metadata") and val:
try:
img_meta[key] = json.loads(val)
except Exception:
img_meta[key] = val
else:
img_meta[key] = val
post_img_meta_path = step_dir / "post_image_meta.json"
with open(post_img_meta_path, "w", encoding="utf-8") as f:
json.dump(img_meta, f, indent=2, ensure_ascii=False)
else:
print(
f"Warning: Image metadata row for {post_image_name} not"
" found in images table."
)
conn.close()
# Write sentinel file to indicate successful chunking
try:
(output_dir / ".chunked").touch()
except Exception as e:
print(f"Warning: Failed to write chunked sentinel file: {e}")
print("Chunking trace completion status: success")
return output_dir
def list_devices(self) -> list[dict]:
"""Dynamically queries the ADB server for connected Android devices."""
try:
from adbutils import AdbClient
adb = AdbClient(host="localhost", port=5037)
devices = adb.device_list()
return [{"serial": d.serial, "status": "online"} for d in devices]
except Exception as e:
print(f"Warning: Failed to query device list from ADB: {e}")
return []
def load_session_goal(
self, session_id: str, step_dir: Path, original_db_path: str = None
) -> str:
"""Loads the session goal from either session.json or the original session db."""
session_json_path = step_dir.parent / "session.json"
if session_json_path.exists():
try:
with open(session_json_path, encoding="utf-8") as f:
data = json.load(f)
if isinstance(data, dict) and "initial_goal" in data and data["initial_goal"]:
return data["initial_goal"]
except Exception as e:
print(f"Warning: Failed to load goal from session.json: {e}")
if original_db_path and Path(original_db_path).exists():
try:
conn = sqlite3.connect(original_db_path)
conn.row_factory = sqlite3.Row
cursor = conn.cursor()
cursor.execute(
"SELECT initial_goal FROM sessions WHERE session_id = ?",
(str(session_id),),
)
row = cursor.fetchone()
conn.close()
if row and row["initial_goal"]:
return row["initial_goal"]
except Exception as e:
print(f"Warning: Failed to load goal from original database: {e}")
raise ValueError(
f"Could not load initial_goal for session {session_id} from"
f" session.json or original database {original_db_path}."
)
def extract_replay_data(
self,
step_dir,
original_db_path=None,
agent_name="explorer",
tool_name="ask_explorer",
):
"""Extracts final answer and chronological tool call sequence from traces.json/db."""
step_json_path = step_dir / "step.json"
if not step_json_path.exists():
raise FileNotFoundError(f"step.json not found in {step_dir}")
with open(step_json_path, encoding="utf-8") as f:
step_data = json.load(f)
step_id = step_data["step_id"]
traces_json_path = step_dir / "traces.json"
traces = []
if traces_json_path.exists():
with open(traces_json_path, encoding="utf-8") as f:
traces = json.load(f)
# If traces.json is empty/missing or lacks agent trace, check original DB
agent_trace = next(
(t for t in traces if t.get("type") == "agent" and t.get("name") == agent_name),
None,
)
if (not agent_trace or not traces) and original_db_path and Path(original_db_path).exists():
conn = sqlite3.connect(original_db_path)
conn.row_factory = sqlite3.Row
cursor = conn.cursor()
cursor.execute(
"SELECT * FROM traces WHERE step_id = ? ORDER BY timestamp ASC",
(step_id,),
)
rows = cursor.fetchall()
traces = []
for r in rows:
td = dict(r)
if td.get("payload"):
try:
td["payload"] = json.loads(td["payload"])
except (ValueError, TypeError):
# Non-JSON payload: keep the raw string.
pass
traces.append(td)
conn.close()
agent_trace = next(
(t for t in traces if t.get("type") == "agent" and t.get("name") == agent_name),
None,
)
if not agent_trace:
print(
f"Warning: {agent_name} agent trace not found. Replay will"
" fallback to direct submit_answer."
)
return (
[],
{
"candidates": [],
"fallback_message": f"No {agent_name} trace found.",
},
traces,
)
payload = agent_trace.get("payload") or {}
result_str = payload.get("result", "{}")
if isinstance(result_str, str):
try:
final_submit_args = json.loads(result_str)
except Exception:
final_submit_args = {
"candidates": [],
"fallback_message": result_str,
}
else:
final_submit_args = result_str
# Build sequence of intermediate tool calls under the agent trace
agent_trace_id = agent_trace["trace_id"]
def get_descendants(pid):
child_list = []
for t in traces:
if t.get("parent_trace_id") == pid:
child_list.append(t)
child_list.extend(get_descendants(t["trace_id"]))
return child_list
descendants = get_descendants(agent_trace_id)
descendants.sort(key=lambda x: x.get("timestamp", 0))
tool_calls_sequence = []
for t in descendants:
if t.get("type") == "tool" and t.get("name") not in (
tool_name,
"submit_answer",
):
t_payload = t.get("payload") or {}
t_args = t_payload.get("args") or {}
tool_calls_sequence.append({"name": t.get("name"), "args": t_args})
# Preload traces up to the point when agent is called
traces_to_preload = []
if agent_trace in traces:
idx = traces.index(agent_trace)
traces_to_preload = traces[:idx]
else:
agent_ts = agent_trace.get("timestamp", 0.0)
traces_to_preload = [t for t in traces if t.get("timestamp", 0.0) < agent_ts]
return tool_calls_sequence, final_submit_args, traces_to_preload
def resolve_step_data_path(self, data: dict, path: str):
"""Resolves a value from a nested dict/list using a dot-separated path (e.g.
'action_taken.0.target_text').
"""
parts = path.split(".")
curr = data
for p in parts:
if isinstance(curr, dict):
curr = curr.get(p)
elif isinstance(curr, list):
try:
idx = int(p)
if idx < len(curr):
curr = curr[idx]
else:
return None
except ValueError:
return None
else:
return None
return curr
def setup_temporary_database(
self,
temp_db_path,
temp_traces_dir,
session_id,
step_data,
pre_image_meta,
initial_goal,
traces_to_preload=None,
post_image_meta=None,
replay_id=None,
):
"""Initializes a sandbox SQLite database with preloaded step and screen metadata."""
if not initial_goal:
raise ValueError("initial_goal must be provided and cannot be empty.")
if temp_db_path.exists():
try:
temp_db_path.unlink()
except OSError:
pass
if temp_traces_dir.exists():
shutil.rmtree(temp_traces_dir)
temp_traces_dir.mkdir(parents=True, exist_ok=True)
try:
from artemis.data_engine.storage import StorageManager
from artemis.data_engine.models import (
SessionMetadata,
ImageRecord,
StepRecord,
TraceRecord,
)
except ImportError:
raise ImportError(
"Failed to import Artemis Data Engine modules. Ensure you are"
" running in the correct virtual environment."
)
storage = StorageManager(temp_db_path, temp_traces_dir)
# Preload session metadata
session_meta = SessionMetadata(
session_id=UUID(session_id) if isinstance(session_id, str) else session_id,
initial_goal=initial_goal,
status="running",
)
storage.create_session(session_meta)
if replay_id:
replay_session_meta = SessionMetadata(
session_id=UUID(replay_id) if isinstance(replay_id, str) else replay_id,
initial_goal=initial_goal,
status="running",
)
storage.create_session(replay_session_meta)
# Preload image metadata
if pre_image_meta:
image_record = ImageRecord(
image_name=pre_image_meta["image_name"],
timestamp=pre_image_meta.get("timestamp", 0.0),
ocr_result=pre_image_meta.get("ocr_result"),
ui_tree=pre_image_meta.get("ui_tree"),
extra_metadata=pre_image_meta.get("extra_metadata", {}),
)
storage.create_image(image_record)
if post_image_meta:
image_record = ImageRecord(
image_name=post_image_meta["image_name"],
timestamp=post_image_meta.get("timestamp", 0.0),
ocr_result=post_image_meta.get("ocr_result"),
ui_tree=post_image_meta.get("ui_tree"),
extra_metadata=post_image_meta.get("extra_metadata", {}),
)
storage.create_image(image_record)
# Preload step metadata
step_record = StepRecord(
step_id=step_data["step_id"],
session_id=step_data["session_id"],
step_number=step_data["step_number"],
timestamp=step_data.get("timestamp", 0.0),
pre_image_name=step_data.get("pre_image_name"),
post_image_name=step_data.get("post_image_name"),
summary=step_data.get("summary"),
action_taken=step_data.get("action_taken"),
operator_raw_thinking=step_data.get("operator_raw_thinking"),
operator_native_thinking=step_data.get("operator_native_thinking"),
last_execution_result=step_data.get("last_execution_result"),
extra_metadata=step_data.get("extra_metadata", {}),
)
storage.create_step(step_record)
# Preload trace history
if traces_to_preload:
for t in traces_to_preload:
try:
payload_val = t.get("payload")
if isinstance(payload_val, str):
try:
payload_val = json.loads(payload_val)
except Exception:
payload_val = {}
elif not isinstance(payload_val, dict):
payload_val = {}
trace_rec = TraceRecord(
trace_id=UUID(t["trace_id"])
if isinstance(t["trace_id"], str)
else t["trace_id"],
session_id=UUID(t["session_id"])
if isinstance(t["session_id"], str)
else t["session_id"],
step_id=UUID(t["step_id"]) if t.get("step_id") else None,
parent_trace_id=UUID(t["parent_trace_id"])
if t.get("parent_trace_id")
else None,
type=t["type"],
name=t["name"],
timestamp=t.get("timestamp", 0.0),
duration=t.get("duration"),
status=t.get("status", "success"),
payload=payload_val,
)
storage.create_trace(trace_rec)
except Exception as ex:
print(f"Warning: Failed to preload trace {t.get('trace_id')}: {ex}")
return storage
def get_preloaded_trace_ids_for_step(self, session_id: str, step_num: int) -> set:
traces_json_path = (
self.test_data_dir / f"{session_id}_chunked" / f"step_{step_num:02d}" / "traces.json"
)
if traces_json_path.exists():
try:
with open(traces_json_path, encoding="utf-8") as f:
traces = json.load(f)
return {str(t.get("trace_id")).lower() for t in traces if t.get("trace_id")}
except (OSError, ValueError) as e:
print(
f"Warning: Failed to read preloaded trace ids from {traces_json_path}: {e}",
file=sys.stderr,
)
return set()
def load_traces_from_db(self, db_path: Path, step_num: int, preloaded_trace_ids: set) -> list:
if not db_path.exists():
return []
try:
conn = sqlite3.connect(db_path)
conn.row_factory = sqlite3.Row
cursor = conn.cursor()
cursor.execute(
"SELECT trace_id, session_id, step_id, parent_trace_id, type,"
" name, status, timestamp, duration, payload FROM traces ORDER"
" BY timestamp ASC"
)
rows = cursor.fetchall()
conn.close()
preloaded_set = {str(tid).lower() for tid in preloaded_trace_ids if tid}
step_traces = []
for r in rows:
td = dict(r)
if td.get("payload"):
try:
td["payload"] = json.loads(td["payload"])
except (ValueError, TypeError):
# Non-JSON payload: normalized to {} just below.
pass
if not isinstance(td.get("payload"), dict):
td["payload"] = {}
tid = str(r["trace_id"]).lower()
if tid in preloaded_set:
td["is_preloaded"] = True
else:
td["is_preloaded"] = False
td["step_number"] = step_num
step_traces.append(td)
return step_traces
except Exception as e:
print(
f"Warning: Failed to load traces from database {db_path}: {e}",
file=sys.stderr,
)
return []
def normalize_trace_payload(self, t: dict) -> dict:
if t.get("type") == "tool":
payload = t.get("payload") or {}
args = payload.get("args") or {}
positional_args = payload.get("positional_args") or []
if (not args or args == {}) and positional_args:
name = t.get("name")
if name == "search_by_coords" and len(positional_args) >= 2:
try:
args = {
"nx": int(positional_args[0]),
"ny": int(positional_args[1]),
}
except ValueError:
args = {
"nx": positional_args[0],
"ny": positional_args[1],
}
elif name == "search_ui" and len(positional_args) >= 1:
args = {"search_query": positional_args[0]}
elif name == "detect_objects" and len(positional_args) >= 1:
args = {"labels": positional_args[0]}
else:
args = {"positional_args": positional_args}
payload["args"] = args
t["payload"] = payload
return t
def load_all_session_traces(
self,
temp_db_path: Path,
preloaded_trace_ids: set,
current_step_number: int,
) -> list:
"""Loads and merges all traces across all steps in the session."""
session_id = None
try:
session_id = temp_db_path.parents[1].name.split("_")[0]
except IndexError:
# Path too shallow to carry a session dir: fall back to the DB probe.
pass
if not session_id:
try:
conn = sqlite3.connect(temp_db_path)
cursor = conn.cursor()
cursor.execute("SELECT session_id FROM traces LIMIT 1")
row = cursor.fetchone()
if row:
session_id = str(row[0])
conn.close()
except sqlite3.Error:
# Unreadable sandbox DB: fall back to single-step trace loading.
pass
if not session_id:
return self.load_traces_from_db(temp_db_path, current_step_number, preloaded_trace_ids)
session_chunked_dir = self.test_data_dir / f"{session_id}_chunked"
all_traces = []
if session_chunked_dir.exists():
step_subdirs = sorted(
[
d
for d in session_chunked_dir.iterdir()
if d.is_dir() and d.name.startswith("step_")
],
key=lambda d: d.name,
)
for subdir in step_subdirs:
try:
step_num = int(subdir.name.split("_")[1])
except (IndexError, ValueError):
continue
step_outputs_dir = self.test_outputs_dir / f"{session_id}_step_{step_num:02d}"
step_db_path = step_outputs_dir / "temp_traces" / "data_engine.db"
step_preloaded_ids = self.get_preloaded_trace_ids_for_step(session_id, step_num)
if step_db_path.exists():
step_traces = self.load_traces_from_db(
step_db_path, step_num, step_preloaded_ids
)
all_traces.extend(step_traces)
else:
traces_json_path = subdir / "traces.json"
if traces_json_path.exists():
try:
with open(traces_json_path, encoding="utf-8") as f:
step_traces = json.load(f)
for t in step_traces:
t["is_preloaded"] = True
t["step_number"] = step_num
if not isinstance(t.get("payload"), dict):
t["payload"] = {}
all_traces.append(t)
except Exception as e:
print(
f"Warning: Failed to load traces from {traces_json_path}: {e}",
file=sys.stderr,
)
else:
all_traces = self.load_traces_from_db(
temp_db_path, current_step_number, preloaded_trace_ids
)
# Align live replay traces to the appropriate preloaded step/operator trace
for step_num in range(1, 100):
current_step_id = None
current_step_dir = self.test_data_dir / f"{session_id}_chunked" / f"step_{step_num:02d}"
if not current_step_dir.exists():
continue
step_json_path = current_step_dir / "step.json"
if step_json_path.exists():
try:
with open(step_json_path, encoding="utf-8") as f:
step_data = json.load(f)
current_step_id = step_data.get("step_id")
except (OSError, ValueError) as e:
print(
f"Warning: Failed to read step id from {step_json_path}: {e}",
file=sys.stderr,
)
current_step_id_str = str(current_step_id).lower() if current_step_id else None
# 1. Find live root trace (no parent_trace_id and not preloaded)
live_root_trace = None
for t in all_traces:
if t.get("step_number") != step_num:
continue
if not t.get("is_preloaded") and not t.get("parent_trace_id"):
live_root_trace = t
break
# 2. Find all preloaded traces for this step
step_preloaded_traces = [
t
for t in all_traces
if t.get("step_number") == step_num
and t.get("is_preloaded")
and str(t.get("step_id") or "").lower() == current_step_id_str
]
# Sort chronologically to identify the final response of the step
step_preloaded_traces.sort(key=lambda x: x.get("timestamp", 0.0))
final_response_trace = step_preloaded_traces[-1] if step_preloaded_traces else None
# 3. Branch immediately before the final response of the step
if live_root_trace and final_response_trace:
# Resolve the tool name from the live root trace name
tool_name = "ask_explorer"
agent_name = live_root_trace.get("name")
for t_name, cfg in REPLAY_TOOLS_CONFIG.items():
if cfg.get("agent_name") == agent_name:
tool_name = t_name
break
import uuid
virtual_id = str(uuid.uuid4())
# Branch from the same parent as the final response
parent_id = final_response_trace.get("parent_trace_id")
# Set timestamp to be immediately before the final response
virtual_timestamp = final_response_trace.get("timestamp", 0.0) - 0.001
live_query = None
live_cf = None
if live_root_trace.get("payload") and live_root_trace["payload"].get("args"):
args = live_root_trace["payload"]["args"]
live_query = (
args.get("q") or args.get("query") or args.get("arg2") or args.get("goal")
)
live_cf = (
args.get("cf") or args.get("context_feedback") or args.get("arg3") or ""
)
if not live_query:
live_query = "Replay Execution"
virtual_tool_trace = {
"trace_id": virtual_id,
"session_id": session_id,
"step_id": current_step_id,
"parent_trace_id": parent_id,
"type": "tool",
"name": tool_name,
"status": live_root_trace.get("status", "success"),
"timestamp": virtual_timestamp,
"duration": live_root_trace.get("duration", 0.0),
"payload": {
"args": {
"query": live_query,
"context_feedback": live_cf,
},
"result": (live_root_trace.get("payload", {}).get("result")),
},
"is_preloaded": False,
"step_number": step_num,
}
# Link the live root trace to our new virtual tool trace
live_root_trace["parent_trace_id"] = virtual_id
all_traces.append(virtual_tool_trace)
all_traces = [self.normalize_trace_payload(t) for t in all_traces]
all_traces.sort(key=lambda t: t.get("timestamp", 0.0))
return all_traces
def enrich_all_session_traces(self, all_traces: list, session_id: str) -> list:
"""Enriches all preloaded traces (across all steps) with file URLs to their corresponding screenshots."""
session_chunked_dir = self.test_data_dir / f"{session_id}_chunked"
step_dirs = {}
if session_chunked_dir.exists():
for subdir in session_chunked_dir.iterdir():
if subdir.is_dir() and subdir.name.startswith("step_"):
step_json_path = subdir / "step.json"
if step_json_path.exists():
try:
with open(step_json_path, encoding="utf-8") as f:
step_data = json.load(f)
sid = step_data.get("step_id")
if sid:
step_dirs[str(sid).lower()] = subdir
except (OSError, ValueError):
# Unreadable step.json: screenshots for this step
# simply are not linked into the trace payloads.
pass
for td in all_traces:
if td.get("is_preloaded"):
step_id = str(td.get("step_id") or "").lower()
if step_id and step_id in step_dirs:
step_dir = step_dirs[step_id]
if td.get("type") == "agent":
pre_jpg = step_dir / "pre.jpg"
if pre_jpg.exists():
td["payload"]["pre_screenshot_path"] = f"file://{pre_jpg.resolve()}"
elif td.get("type") == "action":
post_jpg = step_dir / "post.jpg"
if post_jpg.exists():
td["payload"]["post_screenshot_path"] = f"file://{post_jpg.resolve()}"
return all_traces
def load_steps_metadata(self, session_id: str) -> dict[int, dict]:
"""Loads metadata for all steps in the session from the chunked data."""
session_chunked_dir = self.test_data_dir / f"{session_id}_chunked"
steps = {}
if session_chunked_dir.exists():
for subdir in session_chunked_dir.iterdir():
if subdir.is_dir() and subdir.name.startswith("step_"):
step_json_path = subdir / "step.json"
if step_json_path.exists():
try:
with open(step_json_path, encoding="utf-8") as f:
step_data = json.load(f)
step_num = step_data.get("step_number")
if step_num is not None:
steps[int(step_num)] = step_data
except Exception as e:
print(f"Warning: Failed to load step data from {step_json_path}: {e}")
return steps
def make_file_url(self, val):
if isinstance(val, str) and (val.endswith(".jpg") or val.endswith(".png") or "/" in val):
try:
p = Path(val)
if p.exists():
abs_path = str(p.resolve())
return (
f"file://{abs_path}" if abs_path.startswith("/") else f"file:///{abs_path}"
)
except (OSError, ValueError):
# Not a resolvable filesystem path: return the value verbatim.
pass
return val
def clean_payload_paths(self, obj):
if isinstance(obj, dict):
return {k: self.clean_payload_paths(v) for k, v in obj.items()}
elif isinstance(obj, list):
return [self.clean_payload_paths(v) for v in obj]
elif isinstance(obj, str):
return self.make_file_url(obj)
return obj
def sanitize_nodes(self, nodes_list):
sanitized = []
for node in nodes_list:
node["children"] = self.sanitize_nodes(node["children"])
if node["name"] == "node":
sanitized.extend(node["children"])
else:
sanitized.append(node)
return sanitized
def sort_tree_nodes_chronologically(self, nodes_list: list[dict]) -> list[dict]:
for node in nodes_list:
if "children" in node and node["children"]:
node["children"] = self.sort_tree_nodes_chronologically(node["children"])
return sorted(nodes_list, key=lambda x: x.get("timestamp", 0.0) or 0.0)
def extract_tool_calls_from_db(self, temp_db_path: Path, preloaded_trace_ids: set) -> list:
"""Extracts tool calls executed during the current step replay from the sandbox database."""
preloaded_set = {str(tid).lower() for tid in preloaded_trace_ids if tid}
if not temp_db_path.exists():
return []
try:
conn = sqlite3.connect(temp_db_path)
conn.row_factory = sqlite3.Row
cursor = conn.cursor()
cursor.execute(
"SELECT trace_id, name, status, payload FROM traces WHERE type"
" = 'tool' ORDER BY timestamp ASC"
)
rows = cursor.fetchall()
conn.close()
except Exception as e:
print(
f"Warning: Failed to query tool calls from database: {e}",
file=sys.stderr,
)
return []
tool_calls = []
for r in rows:
tid = str(r["trace_id"]).lower()
if tid not in preloaded_set:
payload = {}
if r["payload"]:
try:
payload = json.loads(r["payload"])
except (ValueError, TypeError):
# Non-JSON payload: treated as an empty tool payload.
pass
response_val = payload.get("result")
if response_val is None:
response_val = payload.get("error")
else:
if isinstance(response_val, str) and (
response_val.startswith("{") or response_val.startswith("[")
):
try:
response_val = json.loads(response_val)
except ValueError:
# Looks like JSON but is not: keep the raw string.
pass
args = payload.get("args", {})
positional_args = payload.get("positional_args", [])
if positional_args:
if r["name"] == "search_by_coords" and len(positional_args) >= 2:
try:
args = {
"nx": int(positional_args[0]),
"ny": int(positional_args[1]),
}
except ValueError:
args = {
"nx": positional_args[0],
"ny": positional_args[1],
}
elif r["name"] == "search_ui" and len(positional_args) >= 1:
args = {"search_query": positional_args[0]}
elif r["name"] == "detect_objects" and len(positional_args) >= 1:
args = {"labels": positional_args[0]}
else:
args = {"positional_args": positional_args}
tool_calls.append(
{
"name": r["name"],
"args": args,
"status": r["status"],
"response": response_val,
}
)
return tool_calls
def extract_state_prepopulation_data(self, step_dir, step_data, pre_image_meta):
"""Dynamically extracts state prepopulation fields from chunked test step files."""
latest_ui_hierarchy = None
if pre_image_meta and "ui_tree" in pre_image_meta:
latest_ui_hierarchy = pre_image_meta["ui_tree"]
structured_decisions = None
traces_json_path = step_dir / "traces.json"
if traces_json_path.exists():
try:
with open(traces_json_path, encoding="utf-8") as f:
traces = json.load(f)
operator_trace = next(
(t for t in traces if t.get("type") == "agent" and t.get("name") == "operator"),
None,
)
if operator_trace:
payload = operator_trace.get("payload") or {}
result_val = payload.get("result")
if isinstance(result_val, str):
try:
res_dict = json.loads(result_val)
structured_decisions = res_dict.get("structured_decisions")
except Exception as e:
print(f"Warning: Failed to parse operator trace result JSON: {e}")
elif isinstance(result_val, dict):
structured_decisions = result_val.get("structured_decisions")
except Exception as e:
print(f"Warning: Failed to read structured_decisions from traces.json: {e}")
pass
if not structured_decisions and step_data:
action_taken = step_data.get("action_taken")
if action_taken:
try:
structured_decisions = json.dumps(action_taken)
except Exception as e:
print(f"Warning: Failed to serialize action_taken fallback to JSON: {e}")
pass
focused_app_info = None
if latest_ui_hierarchy:
system_packages = {
"com.android.systemui",
"android",
"com.google.android.apps.nexuslauncher",
}
packages = []
for node in latest_ui_hierarchy:
if isinstance(node, dict):
pkg = node.get("package")
if pkg and pkg not in system_packages and "launcher" not in pkg.lower():
packages.append(pkg)
if packages:
from collections import Counter
focused_app_info = Counter(packages).most_common(1)[0][0]
device_date = None
timestamp = None
if step_data and "timestamp" in step_data:
timestamp = step_data["timestamp"]
elif pre_image_meta and "timestamp" in pre_image_meta:
timestamp = pre_image_meta["timestamp"]
if timestamp is not None:
from datetime import datetime, UTC
try:
dt = datetime.fromtimestamp(timestamp, UTC)
device_date = dt.strftime("%a %b %d %H:%M:%S UTC %Y")
except Exception as e:
print(
f"Warning: Failed to parse timestamp {timestamp} into device date string: {e}"
)
return (
latest_ui_hierarchy,
structured_decisions,
focused_app_info,
device_date,
)
def _ensure_session_chunked(self, session_id: str) -> Path:
"""Ensures that the session traces are chunked and ready for replay."""
session_chunked_dir = self.test_data_dir / f"{session_id}_chunked"
db_steps_count = 0
if self.db_path.exists():
try:
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
cursor.execute(
"SELECT COUNT(*) FROM steps WHERE session_id = ?",
(session_id,),
)
db_steps_count = cursor.fetchone()[0]
conn.close()
except Exception as e:
print(f"Warning: Failed to query step count from DB: {e}")
chunked_steps_count = 0
if session_chunked_dir.exists():
chunked_steps_count = len(list(session_chunked_dir.glob("step_*")))
# Check if directory does not exist, does not contain the sentinel file, or count of chunked steps is out of sync
needs_chunking = (
not session_chunked_dir.exists()
or not (session_chunked_dir / ".chunked").exists()
or db_steps_count != chunked_steps_count
)
if needs_chunking:
try:
print(
f"Chunked data needs update for session {session_id} (DB"
f" steps: {db_steps_count}, chunked steps:"
f" {chunked_steps_count}). Chunking on-the-fly..."
)
self.chunk_session_traces(session_id, session_chunked_dir)
except Exception as e:
print(f"Failed to chunk session {session_id} on-the-fly: {e}")
return session_chunked_dir
def resolve_master_video_path(self, db_recorded_path_str: str) -> Path | None:
"""Resolves the physical master video recording path on disk, searching for renamed folders."""
db_video_path = Path(db_recorded_path_str)
# 1. Direct absolute path check
if db_video_path.exists():
return db_video_path
# 2. Check if relative to workspace
workspace_rel_path = self.workspace_root / db_video_path
if workspace_rel_path.exists():
return workspace_rel_path.resolve()
# 3. Check if under traces_path
traces_rel_path = self.traces_path / db_video_path.name
if traces_rel_path.exists():
return traces_rel_path.resolve()
# 4. Resolve folder renaming prefix
parent_dir_name = db_video_path.parent.name
filename = db_video_path.name
for search_dir in [
self.traces_path,
self.workspace_root / "artemis-traces",
]:
if parent_dir_name and search_dir.exists():
try:
for child in search_dir.iterdir():
if child.is_dir() and (
child.name.startswith(parent_dir_name) or parent_dir_name in child.name
):
candidate = child / filename
if candidate.exists():
return candidate.resolve()
except Exception as e:
print(f"Warning: Error while scanning {search_dir} for renamed folder: {e}")
# 5. Check if under traces_path or artemis-traces with parent_dir_name
for search_dir in [
self.traces_path,
self.workspace_root / "artemis-traces",
]:
fallback_path = search_dir / parent_dir_name / filename
if fallback_path.exists():
return fallback_path.resolve()
return None
def calculate_virtual_start_time(self, session_id: str, step_number: int) -> float | None:
"""Calculates the virtual start time based on the database's final response timestamp."""
if not self.db_path.exists():
return None
try:
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
# Fetch fallback current step timestamp
cursor.execute(
"SELECT timestamp FROM steps WHERE session_id = ? AND step_number = ?",
(session_id, step_number),
)
row = cursor.fetchone()
step_timestamp = row[0] if (row and row[0] is not None) else None
if step_number <= 1:
# First step starts at session start
cursor.execute(
"SELECT start_time FROM sessions WHERE session_id = ?",
(session_id,),
)
row = cursor.fetchone()
db_session_start_time = row[0] if (row and row[0] is not None) else step_timestamp
conn.close()
return db_session_start_time
# Get step_id of previous step
cursor.execute(
"SELECT step_id FROM steps WHERE session_id = ? AND step_number = ?",
(session_id, step_number - 1),
)
row = cursor.fetchone()
if not row:
conn.close()
return step_timestamp
prev_step_id = row[0]
# Get maximum trace timestamp for that step (final response timestamp)
cursor.execute(
"SELECT MAX(timestamp) FROM traces WHERE step_id = ?",
(prev_step_id,),
)
row = cursor.fetchone()
conn.close()
if row and row[0] is not None:
db_final_response_timestamp = row[0]
return db_final_response_timestamp - 0.001
return step_timestamp
except Exception as e:
print(f"Warning: Failed to calculate virtual start time from database: {e}")
return None
def _preemptive_clip_video(
self,
session_id: str,
step_number: int,
step_data: dict,
temp_db_path: Path,
outputs_dir: Path,
replay_id: str = None,
):
"""Locates the master recording, calculates step time bounds, trims the video and updates the temp DB."""
import subprocess
import uuid
# 1. Query the master database's video_recordings table
if not self.db_path.exists():
print(
f"Warning: Master database not found at {self.db_path}."
" Skipping preemptive video clip."
)
return
try:
conn_master = sqlite3.connect(self.db_path)
cursor = conn_master.cursor()
cursor.execute(
"SELECT video_id, device_id, start_time, end_time,"
" local_video_path FROM video_recordings WHERE session_id = ?",
(session_id,),
)
row = cursor.fetchone()
conn_master.close()
except Exception as e:
print(f"Warning: Failed to query video_recordings from master database: {e}")
return
if not row:
print(
f"No video recording found for session {session_id} in master"
" database. Skipping clip."
)
return
(
video_id,
device_id,
video_start_time,
video_end_time,
db_recorded_path,
) = row
if not db_recorded_path:
print("Warning: Recorded path is empty in video_recordings. Skipping clip.")
return
# 2. Locate master recording on disk (renaming logic)
master_video_path = self.resolve_master_video_path(db_recorded_path)
if not master_video_path:
print(
f"Warning: Could not locate master recording on disk for path: {db_recorded_path}"
)
return
# 3. Calculate step bounds using virtual start time
step_start = self.calculate_virtual_start_time(session_id, step_number)
if step_start is None:
step_start = step_data.get("timestamp")
if step_start is None:
print("Warning: Step start timestamp could not be resolved. Skipping clip.")
return
# Look for next step start
step_dir = self.test_data_dir / f"{session_id}_chunked" / f"step_{step_number:02d}"
next_step_dir = step_dir.parent / f"step_{step_number + 1:02d}"
next_step_json_path = next_step_dir / "step.json"
if next_step_json_path.exists():
try:
with open(next_step_json_path, encoding="utf-8") as f:
next_step_data = json.load(f)
next_step_start = next_step_data.get("timestamp", step_start + 15.0)
except Exception as e:
print(f"Warning: Failed to read next step json: {e}")
next_step_start = step_start + 15.0
else:
next_step_start = step_start + 15.0
# Calculate offsets
start_offset = step_start - video_start_time
# clamp start_offset to be non-negative
start_offset = max(0.0, start_offset)
end_offset = next_step_start - video_start_time
duration = max(0.1, end_offset - start_offset)
trimmed_video_path = outputs_dir / "recording_trimmed.mp4"
# 4. Execute FFmpeg trim
cmd = [
"ffmpeg",
"-y",
"-i",
str(master_video_path),
"-ss",
f"{start_offset:.3f}",
"-t",
f"{duration:.3f}",
"-c:v",
"libx264",
str(trimmed_video_path),
]
print(
f"Trimming video segment from {master_video_path}:"
f" offset={start_offset:.3f}s, duration={duration:.3f}s"
)
try:
res = subprocess.run(cmd, capture_output=True, text=True)
if res.returncode != 0:
print(
f"Warning: ffmpeg failed with exit code {res.returncode}. stderr: {res.stderr}"
)
return
except Exception as e:
print(f"Warning: Failed to execute ffmpeg command: {e}")
return
# 5. Update Sandbox Database
try:
conn_temp = sqlite3.connect(temp_db_path)
conn_temp.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)
)
""")
# Insert original metadata first if not present
conn_temp.execute(
"""
INSERT OR IGNORE INTO video_recordings (video_id, session_id, device_id, start_time, end_time, local_video_path)
VALUES (?, ?, ?, ?, ?, ?)
""",
(
video_id,
session_id,
device_id,
video_start_time,
video_end_time,
db_recorded_path,
),
)
# Update values
conn_temp.execute(
"""
UPDATE video_recordings
SET local_video_path = ?, start_time = ?, end_time = ?
WHERE session_id = ?
""",
(
str(trimmed_video_path),
step_start,
next_step_start,
session_id,
),
)
if replay_id:
# Also copy / update for replay_id
conn_temp.execute(
"""
INSERT OR IGNORE INTO video_recordings (video_id, session_id, device_id, start_time, end_time, local_video_path)
VALUES (?, ?, ?, ?, ?, ?)
""",
(
str(uuid.uuid4()),
replay_id,
device_id,
video_start_time,
video_end_time,
db_recorded_path,
),
)
conn_temp.execute(
"""
UPDATE video_recordings
SET local_video_path = ?, start_time = ?, end_time = ?
WHERE session_id = ?
""",
(
str(trimmed_video_path),
step_start,
next_step_start,
replay_id,
),
)
conn_temp.commit()
conn_temp.close()
print(
f"Successfully updated sandbox database at {temp_db_path} with"
" trimmed video segment."
)
except Exception as e:
print(f"Warning: Failed to update sandbox SQLite database: {e}")
def create_ctx(
self,
session_id: str,
step_number: int,
agent_name: str = "explorer",
replay_id: str = None,
) -> ArtemisContext:
"""Creates a sandboxed ArtemisContext and DataEngine for the given session and step."""
self._ensure_session_chunked(session_id)
step_dir = self.test_data_dir / f"{session_id}_chunked" / f"step_{step_number:02d}"
if not step_dir.exists():
raise FileNotFoundError(f"Step directory not found at {step_dir}")
step_json_path = step_dir / "step.json"
pre_image_meta_path = step_dir / "pre_image_meta.json"
post_image_meta_path = step_dir / "post_image_meta.json"
with open(step_json_path, encoding="utf-8") as f:
step_data = json.load(f)
pre_image_meta = None
if pre_image_meta_path.exists():
with open(pre_image_meta_path, encoding="utf-8") as f:
pre_image_meta = json.load(f)
post_image_meta = None
if post_image_meta_path.exists():
with open(post_image_meta_path, encoding="utf-8") as f:
post_image_meta = json.load(f)
outputs_dir = self.test_outputs_dir / f"{session_id}_step_{step_number:02d}"
outputs_dir.mkdir(parents=True, exist_ok=True)
temp_traces_dir = outputs_dir / "temp_traces"
temp_db_path = temp_traces_dir / "data_engine.db"
_, _, traces_to_preload = self.extract_replay_data(
step_dir, str(self.db_path), agent_name=agent_name
)
initial_goal = self.load_session_goal(session_id, step_dir, str(self.db_path))
self.setup_temporary_database(
temp_db_path,
temp_traces_dir,
session_id,
step_data,
pre_image_meta,
initial_goal,
traces_to_preload,
post_image_meta=post_image_meta,
replay_id=replay_id,
)
try:
self._preemptive_clip_video(
session_id=session_id,
step_number=step_number,
step_data=step_data,
temp_db_path=temp_db_path,
outputs_dir=outputs_dir,
replay_id=replay_id,
)
except Exception as e:
print(f"Warning: Pre-emptive video clipping failed: {e}")
try:
from artemis.context import (
ArtemisContext,
DeviceContext,
DevicePlatform,
ExecutionSetup,
)
from artemis.data_engine.engine import DataEngine
ArtemisContext.model_rebuild()
except ImportError as import_err:
raise ImportError(f"Failed to import Artemis core modules: {import_err}")
import sqlite3
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
cursor.execute(
"SELECT device_info FROM sessions WHERE session_id = ?",
(session_id,),
)
row = cursor.fetchone()
conn.close()
device_info = {}
if row and row[0]:
try:
device_info = json.loads(row[0])
except (ValueError, TypeError):
# Malformed device_info: simulate with the defaults below.
pass
sim_device_id = device_info.get("device_id", "replay-device")
sim_w = device_info.get("device_width", 1080)
sim_h = device_info.get("device_height", 2400)
device_context = DeviceContext(
host_platform="LINUX",
mobile_platform=DevicePlatform.ANDROID,
device_id=sim_device_id,
device_width=sim_w,
device_height=sim_h,
)
print(f"Replay simulated device: {sim_device_id} ({sim_w}x{sim_h})")
from artemis.config import get_default_llm_config
try:
llm_cfg = get_default_llm_config()
except Exception as e:
print(f"Warning: Failed to load default LLM config: {e}")
llm_cfg = None
ctx = ArtemisContext(
device=device_context,
adb_client=None,
ui_adb_client=None,
execution_setup=ExecutionSetup(traces_path=temp_traces_dir),
llm_config=llm_cfg,
)
data_engine = DataEngine(ctx)
data_engine.current_session_id = replay_id or session_id
data_engine.current_step_id = step_data["step_id"]
data_engine.current_step_dir = outputs_dir
ctx.data_engine = data_engine
settings.DATA_ENGINE_DB_PATH = temp_db_path
settings.TRACES_PATH = temp_traces_dir
os.environ["DATA_ENGINE_DB_PATH"] = str(temp_db_path)
return ctx
def instantiate_state(
self,
session_id: str,
step_number: int,
) -> State:
"""Instantiates the State object for the given session and step number."""
self._ensure_session_chunked(session_id)
step_dir = self.test_data_dir / f"{session_id}_chunked" / f"step_{step_number:02d}"
if not step_dir.exists():
raise FileNotFoundError(f"Step directory not found at {step_dir}")
step_json_path = step_dir / "step.json"
pre_image_meta_path = step_dir / "pre_image_meta.json"
with open(step_json_path, encoding="utf-8") as f:
step_data = json.load(f)
pre_image_meta = None
if pre_image_meta_path.exists():
with open(pre_image_meta_path, encoding="utf-8") as f:
pre_image_meta = json.load(f)
initial_goal = self.load_session_goal(session_id, step_dir, str(self.db_path))
ui_hier, decisions, app_info, dev_date = self.extract_state_prepopulation_data(
step_dir, step_data, pre_image_meta
)
try:
from artemis.graph.state import State
except ImportError as import_err:
raise ImportError(f"Failed to import State module: {import_err}")
state = State(
messages=[],
initial_goal=initial_goal,
latest_screenshot=str(step_dir / "pre.jpg")
if (step_dir / "pre.jpg").exists()
else str(step_dir / "post.jpg"),
operator_raw_data={"width": self.w, "height": self.h},
current_step_id=step_data["step_id"],
complete_subgoals_by_ids=[],
validator_messages=[],
subagent_calls=step_data.get("subagent_calls", []),
latest_ui_hierarchy=ui_hier,
focused_app_info=app_info,
device_date=dev_date,
structured_decisions=decisions,
)
return state
def generate_replay_id(self) -> str:
"""Generates a new unique replay ID."""
import uuid
return str(uuid.uuid4())
async def replay_step_tool(
self,
session_id: str,
step_number: int,
override_device_id: str = None,
user_submits: dict = None,
tool_name: str = "ask_explorer",
replay_id: str = None,
) -> list[dict]:
"""Replays a specific step tool in a sandboxed execution environment."""
if not replay_id:
replay_id = self.generate_replay_id()
replay_result = {
"session_id": session_id,
"step_number": step_number,
"replay_id": replay_id,
"success": False,
}
tool_cfg = REPLAY_TOOLS_CONFIG.get(tool_name)
if not tool_cfg:
raise ValueError(f"Unknown tool for replay: {tool_name}")
agent_name = tool_cfg["agent_name"]
outputs_parent = self.test_outputs_dir
if outputs_parent.exists():
for item in outputs_parent.iterdir():
if item.is_dir() and item.name.startswith(f"{session_id}_step_"):
archive_dir = self.replay_base_dir / "data" / "older"
archive_dir.mkdir(parents=True, exist_ok=True)
from datetime import datetime
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
target_name = f"{item.name}_{timestamp}"
target_path = archive_dir / target_name
counter = 1
while target_path.exists():
target_name = f"{item.name}_{timestamp}_{counter}"
target_path = archive_dir / target_name
counter += 1
print(
f"Archiving previous run output directory {item.name} to {target_path}..."
)
try:
shutil.move(str(item), str(target_path))
except Exception as archive_err:
print(
f"Warning: Failed to archive {item.name}: {archive_err}",
file=sys.stderr,
)
outputs_dir = self.test_outputs_dir / f"{session_id}_step_{step_number:02d}"
temp_db_path = outputs_dir / "temp_traces" / "data_engine.db"
# Load step data locally for fallback mapping resolutions
step_dir = self.test_data_dir / f"{session_id}_chunked" / f"step_{step_number:02d}"
step_json_path = step_dir / "step.json"
with open(step_json_path, encoding="utf-8") as f:
step_data = json.load(f)
# Retrieve preloaded trace IDs from chunked trace files to filter new tool calls
_, _, traces_to_preload = self.extract_replay_data(
step_dir, str(self.db_path), agent_name=agent_name
)
preloaded_trace_ids = {t.get("trace_id") for t in traces_to_preload if t.get("trace_id")}
llm_calls_log = []
# Device initialization removed for step replay
ctx = self.create_ctx(
session_id=session_id,
step_number=step_number,
agent_name=agent_name,
replay_id=replay_id,
)
state = self.instantiate_state(
session_id=session_id,
step_number=step_number,
)
data_engine = ctx.data_engine
print("Executing using OpenAI Responses...")
from langchain_openai import ChatOpenAI
from unittest.mock import patch
original_generate = ChatOpenAI._agenerate
async def wrapped_generate(model, messages, *args, **kwargs):
record = {
"prompt": {
"model": model.model_name,
"contents": [
{"role": message.type, "parts": message.content} for message in messages
],
},
"response": None,
"error": None,
}
llm_calls_log.append(record)
try:
result = await original_generate(model, messages, *args, **kwargs)
record["response"] = [
generation.message.model_dump(mode="json") for generation in result.generations
]
return result
except Exception as error:
record["error"] = error.__class__.__name__
raise
try:
try:
import inspect
import importlib
mod = importlib.import_module(tool_cfg["module"])
tool_fn = getattr(mod, tool_cfg["function"])
user_submits = user_submits or {}
sig = inspect.signature(tool_fn)
tool_args = {}
denylist = tool_cfg.get("denylist_args", [])
fallback_mappings = tool_cfg.get("fallback_mappings", {})
user_submits = user_submits or {}
sig = inspect.signature(tool_fn)
tool_args = {}
for name, param in sig.parameters.items():
if name in denylist:
if name == "ctx":
tool_args["ctx"] = ctx
elif name == "state":
tool_args["state"] = state
continue
if name in user_submits:
val = user_submits[name]
if param.annotation is bool and isinstance(val, str):
val = val.lower() == "true"
tool_args[name] = val
else:
val = ""
if param.default is not inspect.Parameter.empty:
val = param.default
if name in fallback_mappings:
fallback_path = fallback_mappings[name]
resolved_val = self.resolve_step_data_path(step_data, fallback_path)
if resolved_val is not None:
val = resolved_val
tool_args[name] = val
with patch.object(ChatOpenAI, "_agenerate", wrapped_generate):
outcome = await tool_fn(**tool_args)
replay_result["success"] = True
replay_result["result"] = outcome
print(f"Replay Execution complete. Outcome: {outcome}")
except Exception as replay_err:
print(f"Error during step replay: {replay_err}", file=sys.stderr)
traceback.print_exc()
replay_result["success"] = False
replay_result["error"] = (
f"{replay_err.__class__.__name__}: {str(replay_err)}"
if str(replay_err)
else replay_err.__class__.__name__
)
raise replay_err
finally:
await data_engine.shutdown()
tool_calls = []
if temp_db_path.exists():
try:
tool_calls = self.extract_tool_calls_from_db(temp_db_path, preloaded_trace_ids)
except Exception as ex:
print(
f"Warning: Failed to extract tool calls: {ex}",
file=sys.stderr,
)
replay_result["tool_calls"] = tool_calls
replay_result["llm_calls"] = llm_calls_log
outcome_path = outputs_dir / "result.json"
try:
with open(outcome_path, "w", encoding="utf-8") as f:
json.dump(replay_result, f, indent=2, ensure_ascii=False)
except Exception as write_err:
print(
f"Warning: Failed to write result.json: {write_err}",
file=sys.stderr,
)
return llm_calls_log
def get_replay_config(self, tool_name: str = "ask_explorer") -> list[dict]:
"""Returns the parameter list and types of the selected tool logic."""
tool_cfg = REPLAY_TOOLS_CONFIG.get(tool_name)
if not tool_cfg:
raise ValueError(f"Unknown tool for replay: {tool_name}")
try:
import inspect
import importlib
mod = importlib.import_module(tool_cfg["module"])
tool_fn = getattr(mod, tool_cfg["function"])
sig = inspect.signature(tool_fn)
denylist = tool_cfg.get("denylist_args", [])
explorer_params = []
for name, param in sig.parameters.items():
if name in denylist:
continue
default_val = ""
if param.default is not inspect.Parameter.empty:
default_val = param.default
input_type = "text"
if any(
keyword in name.lower()
for keyword in (
"feedback",
"context",
"description",
"note",
)
):
input_type = "textarea"
elif param.annotation is bool:
input_type = "checkbox"
elif param.annotation in (int, float):
input_type = "number"
explorer_params.append(
{
"name": name,
"default": default_val,
"input_type": input_type,
}
)
return explorer_params
except Exception as e:
raise Exception(f"Failed to load signature for {tool_name}: {str(e)}")
def get_replay_tools(self) -> list[dict]:
"""Returns the list of registered tools available for replay."""
return [
{
"name": name,
"display_name": cfg["display_name"],
"description": cfg["description"],
}
for name, cfg in REPLAY_TOOLS_CONFIG.items()
]
def get_replay_steps(self, session_id: str) -> list[dict]:
"""Loads and formats metadata for all chunked steps in the session."""
session_chunked_dir = self._ensure_session_chunked(session_id)
if not session_chunked_dir.exists():
raise FileNotFoundError("Step replay chunked data not found for this session.")
try:
steps_metadata = self.load_steps_metadata(session_id)
step_numbers = sorted(list(steps_metadata.keys()))
import inspect
import importlib
# Pre-load signatures and configs for all registered tools
tool_signatures = {}
for t_name, t_cfg in REPLAY_TOOLS_CONFIG.items():
try:
mod = importlib.import_module(t_cfg["module"])
fn = getattr(mod, t_cfg["function"])
tool_signatures[t_name] = inspect.signature(fn)
except Exception as e:
print(
f"Warning: Failed to load signature for tool {t_name}: {e}",
file=sys.stderr,
)
steps_list = []
for num in step_numbers:
s_data = steps_metadata.get(num, {})
sid = s_data.get("step_id")
pre_url = None
post_url = None
step_dir = session_chunked_dir / f"step_{num:02d}"
if step_dir.exists():
pre_jpg = step_dir / "pre.jpg"
if pre_jpg.exists():
pre_url = f"file://{pre_jpg.resolve()}"
post_jpg = step_dir / "post.jpg"
if post_jpg.exists():
post_url = f"file://{post_jpg.resolve()}"
if pre_url and post_url and pre_url == post_url:
post_url = None
# Generate param_defaults for all registered tools
all_tool_defaults = {}
for t_name, sig in tool_signatures.items():
t_cfg = REPLAY_TOOLS_CONFIG[t_name]
denylist = t_cfg.get("denylist_args", [])
fallback_mappings = t_cfg.get("fallback_mappings", {})
defaults = {}
for name, param in sig.parameters.items():
if name in denylist:
continue
val = ""
if param.default is not inspect.Parameter.empty:
val = param.default
# Resolve fallback using mappings
if name in fallback_mappings:
mapping_path = fallback_mappings[name]
parts = mapping_path.split(".")
curr = s_data
for part in parts:
if isinstance(curr, dict):
curr = curr.get(part)
elif isinstance(curr, list):
try:
idx = int(part)
if idx < len(curr):
curr = curr[idx]
else:
curr = None
except ValueError:
curr = None
else:
curr = None
if curr is not None:
val = curr
defaults[name] = val
all_tool_defaults[t_name] = defaults
# Check if this step has already been replayed by scanning outputs folder
step_outputs_dir = self.test_outputs_dir / f"{session_id}_step_{num:02d}"
step_db_path = step_outputs_dir / "temp_traces" / "data_engine.db"
is_replayed = step_db_path.exists()
steps_list.append(
{
"step_id": str(sid).lower() if sid else None,
"step_number": num,
"timestamp": s_data.get("timestamp"),
"summary": s_data.get("summary", "No summary available."),
"action_taken": s_data.get("action_taken", []),
"query": "",
"param_defaults": all_tool_defaults,
"pre_screenshot": pre_url,
"post_screenshot": post_url,
"operator_raw_thinking": s_data.get("operator_raw_thinking"),
"operator_native_thinking": s_data.get("operator_native_thinking"),
"is_replayed": is_replayed,
}
)
return steps_list
except Exception as e:
traceback.print_exc()
raise e
async def run_step_replay(
self,
session_id: str,
step_number: int,
device_id: str,
user_submits: dict,
tool_name: str = "ask_explorer",
replay_id: str = None,
) -> dict:
"""Triggers sandbox execution of step replay and returns the resulting traces trees."""
tool_cfg = REPLAY_TOOLS_CONFIG.get(tool_name)
if not tool_cfg:
raise ValueError(f"Unknown tool for replay: {tool_name}")
agent_name = tool_cfg["agent_name"]
# 1. Run live step replay
llm_calls_log = await self.replay_step_tool(
session_id=session_id,
step_number=step_number,
override_device_id=device_id,
user_submits=user_submits,
tool_name=tool_name,
replay_id=replay_id,
)
# 2. Extract freshly generated traces
outputs_dir = self.test_outputs_dir / f"{session_id}_step_{step_number:02d}"
temp_db_path = outputs_dir / "temp_traces" / "data_engine.db"
self._ensure_session_chunked(session_id)
step_dir = self.test_data_dir / f"{session_id}_chunked" / f"step_{step_number:02d}"
_, _, traces_to_preload = self.extract_replay_data(
step_dir, str(self.db_path), agent_name=agent_name
)
preloaded_trace_ids = {t.get("trace_id") for t in traces_to_preload if t.get("trace_id")}
all_traces = self.load_all_session_traces(temp_db_path, preloaded_trace_ids, step_number)
# Align LLM calls and extract final tool result from result.json if available
replay_outcome = None
replay_id = None
result_json_path = outputs_dir / "result.json"
if result_json_path.exists():
try:
with open(result_json_path, encoding="utf-8") as f:
res_data = json.load(f)
replay_outcome = res_data.get("result")
replay_id = res_data.get("replay_id")
except (OSError, ValueError) as e:
print(
f"Warning: Failed to read replay result from {result_json_path}: {e}",
file=sys.stderr,
)
# Update virtual tool trace result with the actual formatted outcome if available
if replay_outcome is not None:
for node in all_traces:
if (
not node.get("is_preloaded")
and node["type"] == "tool"
and node["name"] == tool_name
):
node["payload"]["result"] = replay_outcome
# Collect all valid LLM span names from REPLAY_TOOLS_CONFIG to align subagent calls too
all_llm_span_names = {
cfg.get("llm_span_name")
for cfg in REPLAY_TOOLS_CONFIG.values()
if cfg.get("llm_span_name")
}
llm_call_index = 0
for node in all_traces:
if (
not node.get("is_preloaded")
and node["type"] == "span"
and node["name"] in all_llm_span_names
and llm_calls_log
and llm_call_index < len(llm_calls_log)
):
node["type"] = "llm_call"
node["llm_call"] = llm_calls_log[llm_call_index]
llm_call_index += 1
# Clean payload paths
for node in all_traces:
if node.get("payload"):
node["payload"] = self.clean_payload_paths(node["payload"])
if node.get("llm_call"):
node["llm_call"] = self.clean_payload_paths(node["llm_call"])
# Filter all_traces for ONLY this step's traces
step_traces = [t for t in all_traces if t.get("step_number") == step_number]
preloaded_traces = [t for t in step_traces if t.get("is_preloaded")]
live_traces = [t for t in step_traces if not t.get("is_preloaded")]
# Build preloaded hierarchy
preloaded_nodes = {
str(t["trace_id"]).lower(): {**t, "children": []} for t in preloaded_traces
}
preloaded_root_nodes = []
for node_id, node in preloaded_nodes.items():
parent_id = node.get("parent_trace_id")
if parent_id:
parent_id_str = str(parent_id).lower()
if parent_id_str in preloaded_nodes:
preloaded_nodes[parent_id_str]["children"].append(node)
else:
preloaded_root_nodes.append(node)
else:
preloaded_root_nodes.append(node)
preloaded_root_nodes = self.sanitize_nodes(preloaded_root_nodes)
preloaded_root_nodes = self.sort_tree_nodes_chronologically(preloaded_root_nodes)
# Build live hierarchy
live_nodes = {str(t["trace_id"]).lower(): {**t, "children": []} for t in live_traces}
# Add a virtual "Response" child node at the end of the virtual tool trace to show the final outcome at the bottom of the tree
if replay_outcome is not None:
virtual_tool_node_id = None
for node_id, node in live_nodes.items():
if node.get("type") == "tool" and node.get("name") == tool_name:
virtual_tool_node_id = node_id
break
if virtual_tool_node_id:
import uuid
outcome_node_id = str(uuid.uuid4())
max_ts = (
max([t.get("timestamp", 0.0) for t in live_traces])
if live_traces
else live_nodes[virtual_tool_node_id].get("timestamp", 0.0)
)
outcome_node = {
"trace_id": outcome_node_id,
"session_id": session_id,
"step_id": live_nodes[virtual_tool_node_id].get("step_id"),
"parent_trace_id": live_nodes[virtual_tool_node_id]["trace_id"],
"type": "tool",
"name": f"{tool_name} Response",
"status": (live_nodes[virtual_tool_node_id].get("status", "success")),
"timestamp": max_ts + 0.001,
"duration": 0.0,
"payload": {"result": replay_outcome},
"is_preloaded": False,
"step_number": step_number,
"children": [],
}
live_nodes[outcome_node_id] = outcome_node
live_root_nodes = []
for node_id, node in live_nodes.items():
parent_id = node.get("parent_trace_id")
if parent_id:
parent_id_str = str(parent_id).lower()
if parent_id_str in live_nodes:
live_nodes[parent_id_str]["children"].append(node)
else:
live_root_nodes.append(node)
else:
live_root_nodes.append(node)
live_root_nodes = self.sanitize_nodes(live_root_nodes)
live_root_nodes = self.sort_tree_nodes_chronologically(live_root_nodes)
return {
"success": True,
"preloaded": preloaded_root_nodes,
"live": live_root_nodes,
"replay_id": replay_id,
}
def get_step_replay_traces(
self, session_id: str, step_number: int, tool_name: str = "ask_explorer"
) -> dict:
"""Retrieves previously generated step replay traces if they exist."""
tool_cfg = REPLAY_TOOLS_CONFIG.get(tool_name)
if not tool_cfg:
raise ValueError(f"Unknown tool for replay: {tool_name}")
agent_name = tool_cfg["agent_name"]
outputs_dir = self.test_outputs_dir / f"{session_id}_step_{step_number:02d}"
temp_db_path = outputs_dir / "temp_traces" / "data_engine.db"
if not temp_db_path.exists():
return {"success": True, "live": []}
self._ensure_session_chunked(session_id)
step_dir = self.test_data_dir / f"{session_id}_chunked" / f"step_{step_number:02d}"
preloaded_trace_ids = set()
if step_dir.exists():
_, _, traces_to_preload = self.extract_replay_data(
step_dir,
str(self.db_path),
agent_name=agent_name,
tool_name=tool_name,
)
preloaded_trace_ids = {
t.get("trace_id") for t in traces_to_preload if t.get("trace_id")
}
all_traces = self.load_all_session_traces(temp_db_path, preloaded_trace_ids, step_number)
llm_span_name = tool_cfg.get("llm_span_name")
if not llm_span_name:
raise ValueError(f"Tool {tool_name} config is missing 'llm_span_name'")
# Align LLM calls and extract final tool result from result.json if available
llm_calls_log = []
replay_outcome = None
replay_id = None
result_json_path = outputs_dir / "result.json"
if result_json_path.exists():
try:
with open(result_json_path, encoding="utf-8") as f:
res_data = json.load(f)
llm_calls_log = res_data.get("llm_calls") or []
replay_outcome = res_data.get("result")
replay_id = res_data.get("replay_id")
except (OSError, ValueError) as e:
print(
f"Warning: Failed to read replay result from {result_json_path}: {e}",
file=sys.stderr,
)
# Update virtual tool trace result with the actual formatted outcome if available
if replay_outcome is not None:
for node in all_traces:
if (
not node.get("is_preloaded")
and node["type"] == "tool"
and node["name"] == tool_name
):
node["payload"]["result"] = replay_outcome
# Collect all valid LLM span names from REPLAY_TOOLS_CONFIG to align subagent calls too
all_llm_span_names = {
cfg.get("llm_span_name")
for cfg in REPLAY_TOOLS_CONFIG.values()
if cfg.get("llm_span_name")
}
llm_call_index = 0
for node in all_traces:
if (
not node.get("is_preloaded")
and node["type"] == "span"
and node["name"] in all_llm_span_names
and llm_calls_log
and llm_call_index < len(llm_calls_log)
):
node["type"] = "llm_call"
node["llm_call"] = llm_calls_log[llm_call_index]
llm_call_index += 1
# Clean payload paths
for node in all_traces:
if node.get("payload"):
node["payload"] = self.clean_payload_paths(node["payload"])
if node.get("llm_call"):
node["llm_call"] = self.clean_payload_paths(node["llm_call"])
# Filter all_traces for ONLY this step's traces
step_traces = [t for t in all_traces if t.get("step_number") == step_number]
live_traces = [t for t in step_traces if not t.get("is_preloaded")]
# Build live hierarchy
live_nodes = {str(t["trace_id"]).lower(): {**t, "children": []} for t in live_traces}
# Add a virtual "Response" child node at the end of the virtual tool trace to show the final outcome at the bottom of the tree
if replay_outcome is not None:
virtual_tool_node_id = None
for node_id, node in live_nodes.items():
if node.get("type") == "tool" and node.get("name") == tool_name:
virtual_tool_node_id = node_id
break
if virtual_tool_node_id:
import uuid
outcome_node_id = str(uuid.uuid4())
max_ts = (
max([t.get("timestamp", 0.0) for t in live_traces])
if live_traces
else live_nodes[virtual_tool_node_id].get("timestamp", 0.0)
)
outcome_node = {
"trace_id": outcome_node_id,
"session_id": session_id,
"step_id": live_nodes[virtual_tool_node_id].get("step_id"),
"parent_trace_id": live_nodes[virtual_tool_node_id]["trace_id"],
"type": "tool",
"name": f"{tool_name} Response",
"status": (live_nodes[virtual_tool_node_id].get("status", "success")),
"timestamp": max_ts + 0.001,
"duration": 0.0,
"payload": {"result": replay_outcome},
"is_preloaded": False,
"step_number": step_number,
"children": [],
}
live_nodes[outcome_node_id] = outcome_node
live_root_nodes = []
for node_id, node in live_nodes.items():
parent_id = node.get("parent_trace_id")
if parent_id:
parent_id_str = str(parent_id).lower()
if parent_id_str in live_nodes:
live_nodes[parent_id_str]["children"].append(node)
else:
live_root_nodes.append(node)
else:
live_root_nodes.append(node)
live_root_nodes = self.sanitize_nodes(live_root_nodes)
live_root_nodes = self.sort_tree_nodes_chronologically(live_root_nodes)
return {
"success": True,
"live": live_root_nodes,
"replay_id": replay_id,
}
Reference in new issue View Git Blame Copy Permalink