diff --git a/apps/admin_console/routers/tasks.py b/apps/admin_console/routers/tasks.py index 38650d6..26fca46 100644 --- a/apps/admin_console/routers/tasks.py +++ b/apps/admin_console/routers/tasks.py @@ -228,6 +228,7 @@ async def stop_task( # Empty or non-JSON body: fall back to the query parameters. pass + # stop_tasks updates scheduler state and asyncio events owned by this loop. stopped = task_queue_service.stop_tasks( clear_all=target_all, session_id=target_sid, diff --git a/apps/admin_console/server.py b/apps/admin_console/server.py index f6f12a2..151bb82 100644 --- a/apps/admin_console/server.py +++ b/apps/admin_console/server.py @@ -19,6 +19,7 @@ Modular entrypoint for full trace inspection, step replay, and task execution ma import argparse import asyncio +from contextlib import asynccontextmanager import os from pathlib import Path import secrets @@ -95,8 +96,19 @@ except ImportError: # Initialize language server synchronization address init_ls_address() + +@asynccontextmanager +async def _lifespan(_app: "FastAPI"): + """Start server resources and release them on shutdown.""" + await on_startup() + try: + yield + finally: + await on_shutdown() + + # Initialize FastAPI application -app = FastAPI(title="Artemis Admin & Trace Console") +app = FastAPI(title="Artemis Admin & Trace Console", lifespan=_lifespan) LIFECYCLE_TOKEN = os.environ.get("ARTEMIS_LIFECYCLE_TOKEN") or secrets.token_urlsafe(32) app.state.lifecycle_token = LIFECYCLE_TOKEN @@ -106,7 +118,6 @@ app.state.lifecycle_token = LIFECYCLE_TOKEN app.add_middleware(SameOriginBoundaryMiddleware) -@app.on_event("startup") async def on_startup(): """Startup lifecycle hooks.""" state.is_shutting_down = False @@ -150,7 +161,6 @@ async def on_startup(): state.worker_task = asyncio.create_task(task_queue_service.queue_worker()) -@app.on_event("shutdown") async def on_shutdown(): """Stop task and IPC children before the UI server exits.""" state.is_shutting_down = True diff --git a/apps/admin_console/services/task_queue_service.py b/apps/admin_console/services/task_queue_service.py index b107bf6..b99b4c9 100644 --- a/apps/admin_console/services/task_queue_service.py +++ b/apps/admin_console/services/task_queue_service.py @@ -343,40 +343,23 @@ class TaskQueueService: pass @classmethod - def _concurrency_mode(cls) -> str: - """Resolve the scheduler concurrency mode. - - Mirrors DeviceExecutionLock's contract (ARTEMIS_CONCURRENCY_MODE / - ARTEMIS_MAX_CONCURRENT_TASKS): - - "global": strict serial -- one task at a time across all devices. - - "per_device": one task per device -- distinct devices run concurrently. - Defaults to per-device concurrency, matching the lock layer's default. - """ - mode = os.environ.get("ARTEMIS_CONCURRENCY_MODE", "").strip().lower() - if mode in ("global", "serial", "1"): - return "global" - if mode in ("per_device", "device", "parallel", "0"): - return "per_device" - try: - if int(os.environ.get("ARTEMIS_MAX_CONCURRENT_TASKS", 0) or 0) == 1: - return "global" - except (TypeError, ValueError): - pass - return "per_device" + def _concurrency_limit(cls) -> int: + """Return 0 for per-device concurrency, or the global task limit.""" + return DeviceExecutionLock.resolve_env_concurrency() @classmethod async def queue_worker(cls): """Persistent dispatcher scheduling pending tasks onto devices. Scans the queue and launches each eligible task as an independent - coroutine. Admission is governed by _concurrency_mode(): "global" admits - one task at a time overall, "per_device" admits one task per device so - distinct devices execute concurrently. Per-device FIFO ordering and + coroutine. Admission is governed by _concurrency_limit(): 0 admits one + task per device so distinct devices execute concurrently, N>=1 admits + at most N tasks across all devices. Per-device FIFO ordering and cross-process mutual exclusion remain enforced by DeviceExecutionLock inside each worker process. """ print( - f"[QueueWorker] Dispatcher initialized (concurrency mode: {cls._concurrency_mode()})." + f"[QueueWorker] Dispatcher initialized (concurrency limit: {cls._concurrency_limit()})." ) try: DeviceExecutionLock.cleanup_stale_locks() @@ -404,10 +387,10 @@ class TaskQueueService: @classmethod def _dispatch_pending_tasks(cls) -> None: - """Launch every pending task admissible under the current concurrency mode.""" - mode = cls._concurrency_mode() + """Launch every pending task admissible under the current concurrency limit.""" + limit = cls._concurrency_limit() state.prune_finished_runs() - if mode == "global" and state.is_running: + if limit == 1 and state.is_running: return # Dispatched-but-not-yet-spawned runs are only visible as queue items in @@ -416,6 +399,20 @@ class TaskQueueService: in_flight = [ i for i in state.queue_items if isinstance(i, dict) and i.get("status") == "running" ] + capacity = None + if limit >= 1: + # Registered workers still have running queue rows until finalization. + active_pids = { + getattr(run.get("process"), "pid", None) for run in state.active_runs.values() + } - {None} + starting = sum( + str(item.get("session_id")) not in state.active_runs + and item.get("pid") not in active_pids + for item in in_flight + ) + capacity = limit - len(state.active_runs) - starting + if capacity <= 0: + return busy_devices = state.busy_device_ids | { cls._task_target(i).lock_key for i in in_flight if i.get("device_serial") } @@ -431,7 +428,7 @@ class TaskQueueService: device = item.get("device_serial") target = cls._task_target(item) - if mode == "per_device": + if limit == 0: # A task without a resolved device may bind to any serial, so it # only launches on an otherwise idle scheduler; the device lock # then allocates freely without contending against active runs. @@ -439,6 +436,9 @@ class TaskQueueService: continue if device is not None and target.lock_key in busy_devices: continue + elif limit > 1 and device is not None and target.lock_key in busy_devices: + # A second worker for this device would wait on its lock. + continue item["status"] = "running" dispatched_any = True @@ -449,8 +449,10 @@ class TaskQueueService: # tasks, and a collected run would strand its queue item forever. cls._run_tasks.add(run_task) run_task.add_done_callback(cls._run_tasks.discard) - if mode == "global": - break + if capacity is not None: + capacity -= 1 + if capacity <= 0: + break @classmethod def _begin_task_run( @@ -584,10 +586,8 @@ class TaskQueueService: } if sess_id: try: - st = trace_store.read_status(str(sess_id)) - if st: - st["pid"] = proc.pid - trace_store.write_status(str(sess_id), st) + # Preserve status updates written concurrently by the worker. + trace_store.update_trace_pid(str(sess_id), proc.pid) except OSError as exc: # Status probes fall back to DB PID bookkeeping, but note it. print( @@ -595,15 +595,22 @@ class TaskQueueService: ) cls._broadcast_startup_progress(sess_id, "process_ready", "Execution process started") ingress_type = str(task_item.get("ingress", "frontend")) - DeviceExecutionLock.transfer_reservation( - str(task_item.get("queue_ticket")), - proc.pid, - description=f"{ingress_type} task: {goal[:120]}", - device_id=device_serial or "pending", - session_id=str(sess_id) if sess_id else None, - ingress=ingress_type, - lock_scope=target.lock_scope, - ) + queue_ticket = task_item.get("queue_ticket") + if queue_ticket: + transferred = DeviceExecutionLock.transfer_reservation( + str(queue_ticket), + proc.pid, + description=f"{ingress_type} task: {goal[:120]}", + device_id=device_serial or "pending", + session_id=str(sess_id) if sess_id else None, + ingress=ingress_type, + lock_scope=target.lock_scope, + ) + if not transferred: + print( + f"[QueueWorker] Could not transfer queue ticket {queue_ticket} to worker" + f" pid {proc.pid}; the worker will queue a fresh ticket itself." + ) @classmethod def _start_output_forwarder( @@ -876,8 +883,19 @@ class TaskQueueService: if proc is not None and proc.returncode is None: await cls._terminate_worker_process(proc) raise - except Exception as e: - print(f"[QueueWorker] Unexpected error executing task [{sess_id}]: {e}") + except Exception: + logger.exception(f"[QueueWorker] Unexpected error executing task [{sess_id}]") + # A spawn failure must also leave a terminal session status. + if sess_id: + try: + new_status = await cls._persist_terminal_session_status( + sess_id, returncode=1, manual_stop=False + ) + cls._announce_session_end(task_item, sess_id, goal or "", new_status, False) + except (OSError, RuntimeError, ValueError): + logger.exception( + f"[QueueWorker] Could not persist failure status for [{sess_id}]" + ) finally: await cls._finish_output_forwarder(output_task) # 5. Clean up the finished task and release this run's scheduling slot @@ -989,14 +1007,8 @@ class TaskQueueService: ) -> dict[str, Any]: """Reserve a device slot and build one pending queue item for a goal.""" sess_id = single_session_id if single_session_id else str(uuid.uuid4()) + # enqueue_tasks resolves the device before creating queue items. assigned_serial = device_serial - if not assigned_serial: - try: - from artemis.runtime import device_pool - - assigned_serial = device_pool.select_device() - except Exception: - assigned_serial = None queue_ticket = DeviceExecutionLock.reserve( description=f"{ingress} task: {goal[:120]}", @@ -1068,6 +1080,14 @@ class TaskQueueService: return rejection_response single_session_id = session_id if (session_id and len(goals) == 1) else None + if not device_serial: + # Device enumeration may block on ADB. + from artemis.runtime import device_pool + + try: + device_serial = await device_pool.select_device_async() + except Exception: + device_serial = None for i, goal in enumerate(goals): task_item = cls._create_queue_item( goal, @@ -1426,10 +1446,8 @@ class TaskQueueService: global_queued = DeviceExecutionLock.get_queued_tasks() for q_item in global_queued: if str(q_item.get("session_id")) == target_sid: - ticket = q_item.get("queue_ticket") - if ticket: - DeviceExecutionLock.cancel_reservation(ticket) - stopped = True + # get_queued_tasks exposes the reservation as "token". + stopped = DeviceExecutionLock.cancel_reservation(q_item.get("token")) break # Fallback 2: check session repository for a running session with a live worker PID if not stopped: diff --git a/artemis/__init__.py b/artemis/__init__.py index 2c9df42..658bd19 100644 --- a/artemis/__init__.py +++ b/artemis/__init__.py @@ -22,7 +22,6 @@ from artemis_client import Capabilities, Device, TaskHandle, TaskResult from artemis.sdk.agent import Agent from artemis.sdk.builders import Builders from artemis.drivers.base import BaseDeviceDriver -from artemis.tools.base import artemis_tool, ToolRegistry __all__ = [ "__version__", @@ -38,6 +37,4 @@ __all__ = [ "TaskHandle", "TaskResult", "BaseDeviceDriver", - "artemis_tool", - "ToolRegistry", ] diff --git a/artemis/agents/diagnoser/diagnoser.py b/artemis/agents/diagnoser/diagnoser.py index dc8a71c..a27cb60 100644 --- a/artemis/agents/diagnoser/diagnoser.py +++ b/artemis/agents/diagnoser/diagnoser.py @@ -20,7 +20,7 @@ from pathlib import Path from typing import Any import uuid -from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage +from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage, ToolMessage from artemis.context import ArtemisContext from artemis.data_engine.trace import ( @@ -52,6 +52,12 @@ from artemis.utils.task_tree import ( logger = get_logger(__name__) +_MAX_ITERATIONS = 30 + +_FINAL_TURN_WARNING = ( + "[WARNING] This is your final iteration. Call 'submit_answer' to provide your final diagnosis." +) + class Diagnoser: def __init__(self, ctx: ArtemisContext): @@ -69,11 +75,8 @@ class Diagnoser: 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 + def _load_prompt_config(self) -> tuple[str, str]: + """Loads the system prompt and checker-feedback template from disk.""" prompt_path = Path(__file__).parent.joinpath("diagnoser.json") if not prompt_path.exists(): raise FileNotFoundError(f"Prompt file not found at {prompt_path}") @@ -87,27 +90,18 @@ class Diagnoser: checker_feedback_template = prompt_data.get("checker_feedback_template", "") if isinstance(checker_feedback_template, list): checker_feedback_template = "\n".join(checker_feedback_template) + return system_prompt, 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) - + def _build_traced_tools(self) -> list: + """Assembles the tool set (honoring the offline/video gates) and wraps it in tracing.""" all_tools = [ - log_tool, - read_note_tool, - list_notes_tool, - adb_command_tool, - wait_tool, - ui_hierarchy_tool, - submit_answer_tool, + get_analyze_logs_tool(self.ctx), + get_read_note_tool(self.ctx), + get_list_notes_tool(self.ctx), + get_run_short_adb_command_tool(self.ctx), + get_wait_tool(self.ctx), + get_ui_hierarchy_tool(self.ctx), + get_submit_answer_tool(self.ctx), ] if not self.is_device_online: logger.info( @@ -123,17 +117,12 @@ class Diagnoser: ] 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) + all_tools.append(get_video_analyzer_tool(self.ctx, role="diagnoser")) - 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 + return [trace_langchain_tool(t, self.ctx) for t in all_tools] + def _build_plan_and_history(self) -> str: + """Renders the task plan and recent step history from the DataEngine.""" task_plan = "No task plan yet." steps = [] subgoal_hash = "default" @@ -153,7 +142,6 @@ class Diagnoser: 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: @@ -174,30 +162,38 @@ class Diagnoser: 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}" + return plan_and_history - # Get verification findings injected by the checkpoint harvest, if any - checker_feedback = "" + def _collect_checker_feedback(self, state: State) -> str: + """Renders verification findings injected by the checkpoint harvest, if any.""" findings = getattr(state, "operator_feedback", None) if findings: - checker_feedback = "\n".join(f"- {f}" for f in findings) - - # Get Current Screenshot and UI Hierarchy - latest_screenshot_b64 = None - minimal_list = "" + return "\n".join(f"- {f}" for f in findings) + return "" + def _read_screenshot_b64(self, state: State) -> str | None: + """Reads the latest screenshot as base64, or ``None`` when unavailable.""" 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") + return base64.b64encode(f.read()).decode("utf-8") except Exception as e: logger.error(f"Failed to read latest screenshot in Diagnoser: {e}") + return None - if state.latest_ui_hierarchy: - minimal_list = self._format_minimal_list(state.latest_ui_hierarchy) - - # Construct rich human message + def _build_initial_messages( + self, + prompt: str, + state: State, + system_prompt: str, + checker_feedback_template: str, + plan_and_history: str, + checker_feedback: str, + minimal_list: str, + ) -> list[BaseMessage]: + """Build the initial prompt with history, device state and screenshot.""" content = [] - content.append({"type": "text", "text": f"Initial Goal: {initial_goal}\n"}) + content.append({"type": "text", "text": f"Initial Goal: {state.initial_goal}\n"}) content.append( { "type": "text", @@ -217,6 +213,7 @@ class Diagnoser: } ) + latest_screenshot_b64 = self._read_screenshot_b64(state) if latest_screenshot_b64: content.append({"type": "text", "text": "--- Current Screenshot ---"}) content.append( @@ -228,117 +225,358 @@ class Diagnoser: content.append({"type": "text", "text": f"--- Diagnostic Query ---\n{prompt}"}) - current_messages = [ + return [ SystemMessage(content=system_prompt), HumanMessage(content=content), ] - self.ctx.background_jobs = {} + async def _await_running_jobs(self) -> None: + """Blocks until every still-running background job has settled.""" + running_jobs = [ + job for job in self.ctx.background_jobs.values() if job.get("status") == "running" + ] + if running_jobs: + await asyncio.gather( + *(job["task"] for job in running_jobs if job.get("task") is not None), + return_exceptions=True, + ) - max_iterations = 30 + def _drain_finished_jobs(self) -> list[ToolMessage]: + """Converts finished, unconsumed background jobs into tool messages.""" + 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 + return new_tool_messages + + def _cancel_running_jobs(self) -> None: + """Cancels any background job still running when the loop exits.""" + for job in self.ctx.background_jobs.values(): + if job.get("status") == "running" and job.get("task") is not None: + job["task"].cancel() + + def _prepare_iteration_llm( + self, + base_llm, + traced_tools: list, + iteration: int, + current_messages: list[BaseMessage], + ): + """Appends the iteration notice and binds the tools available this turn. + + The final iteration is restricted to ``submit_answer`` so the model + must conclude; earlier iterations get a progress nudge and the full + tool set. + """ + if iteration == _MAX_ITERATIONS - 1: + current_messages.append(HumanMessage(content=_FINAL_TURN_WARNING)) + submit_tool = next( + (t for t in traced_tools if t.name == "submit_answer"), + None, + ) + return base_llm.bind_tools(tools=[submit_tool] if submit_tool else []) + + if iteration > 0: + current_messages.append( + HumanMessage( + content=( + "You have not completed the diagnosis yet" + f" (iteration {iteration + 1} of {_MAX_ITERATIONS})." + ) + ) + ) + return base_llm.bind_tools(tools=traced_tools) + + def _token_usage_payload(self, response) -> dict | None: + """Extracts token-usage counters from the response metadata, if present.""" + if not hasattr(response, "usage_metadata") or response.usage_metadata is None: + return 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) + + return { + "prompt_token_count": prompt_token_count, + "candidates_token_count": candidates_token_count, + "cached_content_token_count": (cached_content_token_count), + } + + async def _invoke_model( + self, + llm, + current_messages: list[BaseMessage], + context_components: dict, + ): + """Invokes the model inside a trace span carrying the context payload.""" + with TraceSpan(name="gemini_diagnoser_call", ctx=self.ctx) as span: + span.payload = {"context_components": context_components} + response = await invoke_llm_with_timeout_message(acomplete(llm, current_messages)) + span.result = response.content if hasattr(response, "content") else str(response) + + token_usage = self._token_usage_payload(response) + if token_usage is not None: + if span.payload is None: + span.payload = {} + span.payload["token_usage"] = token_usage + return response + + def _find_submit_call(self, tool_calls: list) -> dict | None: + """Finds a ``submit_answer`` call among the response's tool calls.""" + return next( + ( + tc + for tc in tool_calls + if tc["name"] == "submit_answer" or tc["name"].endswith(":submit_answer") + ), + None, + ) + + def _handle_submit_call(self, submit_call: dict, traced_tools: list) -> str: + """Records the submitted answer and renders it as the run's outcome.""" + 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} | Actionable Steps: " + for step in steps: + outcome += f"{step} " + return outcome + + def _max_iteration_outcome(self, response, outcome: str) -> str: + """Renders the outcome when the loop cap is hit with tool calls pending.""" + 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": + return ( + "Diagnoser reached the maximum iteration limit of" + f" {_MAX_ITERATIONS} while attempting to call tools:" + f" {tool_names}" + ) + return outcome + warning + + async def _run_video_job(self, job_id: str, tc: dict, traced_tools: list, state: State) -> None: + """Runs a ``video_analyzer`` call in the background and records its result.""" + 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") + + def _start_video_job(self, tc: dict, traced_tools: list, state: State) -> ToolMessage: + """Registers a background job for ``video_analyzer`` and acknowledges it.""" + 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']}"), + ) + + task = asyncio.create_task(self._run_video_job(job_id, tc, traced_tools, state)) + 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", + ) + + async def _run_tool(self, tc: dict, traced_tools: list, state: State) -> ToolMessage: + """Executes one requested tool call and wraps its result as a tool message.""" + 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": + return self._start_video_job(tc, traced_tools, state) + + 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, + ) + + async def _execute_tool_calls( + self, tool_calls: list, traced_tools: list, state: State + ) -> list[ToolMessage]: + """Executes all requested tools in parallel, skipping ``google_search``.""" + active_tool_calls = [ + tc + for tc in tool_calls + if (tc["name"].split(":")[-1] if ":" in tc["name"] else tc["name"]) != "google_search" + ] + return list( + await asyncio.gather( + *(self._run_tool(tc, traced_tools, state) for tc in active_tool_calls) + ) + ) + + @trace(type="agent", name="diagnoser") + async def run(self, prompt: str, state: State) -> str: + """Runs the Diagnoser and returns its summary.""" + + system_prompt, checker_feedback_template = self._load_prompt_config() + + base_llm = get_llm(ctx=self.ctx, name="diagnoser") + traced_tools = self._build_traced_tools() + + plan_and_history = self._build_plan_and_history() + checker_feedback = self._collect_checker_feedback(state) + minimal_list = ( + self._format_minimal_list(state.latest_ui_hierarchy) + if state.latest_ui_hierarchy + else "" + ) + + current_messages = self._build_initial_messages( + prompt, + state, + system_prompt, + checker_feedback_template, + plan_and_history, + checker_feedback, + minimal_list, + ) + context_components = { + "system_prompt": system_prompt, + "plan_and_history": plan_and_history, + "checker_feedback": checker_feedback, + "visible_ui_elements": minimal_list, + "diagnostic_query": prompt, + } + + self.ctx.background_jobs = {} 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 + for i in range(_MAX_ITERATIONS): + # Include background results before asking for the final answer. + if i == _MAX_ITERATIONS - 1: + await self._await_running_jobs() + new_tool_messages = self._drain_finished_jobs() 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) + llm = self._prepare_iteration_llm(base_llm, traced_tools, i, current_messages) - 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( - acomplete(llm, current_messages) - ) - 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), - } + response = await self._invoke_model(llm, current_messages, context_components) if response.content: outcome = response.content @@ -346,197 +584,21 @@ class Diagnoser: 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, - ) + submit_call = self._find_submit_call(response.tool_calls) 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} " + outcome = self._handle_submit_call(submit_call, traced_tools) 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 + if i == _MAX_ITERATIONS - 1: + outcome = self._max_iteration_outcome(response, outcome) 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) + current_messages.extend( + await self._execute_tool_calls(response.tool_calls, traced_tools, state) + ) 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() + self._cancel_running_jobs() return outcome diff --git a/artemis/agents/explorer/_facade.py b/artemis/agents/explorer/_facade.py deleted file mode 100644 index 8e8063d..0000000 --- a/artemis/agents/explorer/_facade.py +++ /dev/null @@ -1,34 +0,0 @@ -# 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. - -"""Late-bound access to the ``artemis.agents.explorer.explorer`` facade module. - -The Explorer implementation is split across several modules, but external -callers (and ``unittest.mock.patch`` targets in the test suite) address -collaborators such as ``settings``, ``StorageManager``, ``draw_dots``, -``get_llm`` etc. through the historical module path -``artemis.agents.explorer.explorer``. Split-out code therefore resolves those -names at call time via :func:`facade` instead of importing them directly, so -patches applied to the facade module keep affecting the split-out code. - -The import happens inside the function (not at module top level) to stay safe -regardless of which module of the package is imported first. -""" - - -def facade(): - """Return the ``artemis.agents.explorer.explorer`` module object.""" - import artemis.agents.explorer.explorer as explorer_module - - return explorer_module diff --git a/artemis/agents/explorer/constants.py b/artemis/agents/explorer/constants.py index a8ae1e7..872ecef 100644 --- a/artemis/agents/explorer/constants.py +++ b/artemis/agents/explorer/constants.py @@ -22,6 +22,9 @@ Explorer tier the user configured. Tier behavior lives in ASK_EXPLORER_TOOL_NAME = "ask_explorer" +#: Model used when the LLM configuration carries no Explorer entry. +DEFAULT_EXPLORER_MODEL = "gemini-3.8-flash" + ASK_EXPLORER_DESCRIPTION = ( "[EXPLORER] Ask the UI Explorer to locate elements on the current screen" " that are missing from the indexed element list or whose listed" diff --git a/artemis/agents/explorer/explorer.py b/artemis/agents/explorer/explorer.py index f0f1ef8..a302ac9 100644 --- a/artemis/agents/explorer/explorer.py +++ b/artemis/agents/explorer/explorer.py @@ -12,26 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Explorer agent facade. +"""Explorer agent entry point. -This module is the stable entry point (and ``mock.patch`` target namespace) -for the Explorer agent. The implementation is split across sibling modules: +The :class:`Explorer` class combines the following modules: -- ``tiers``: the flash / pro / ultra tier table (single source of truth). +- ``tiers``: engine, tools and budgets for flash / pro / ultra. - ``tool_declarations``: static tool schemas (universal + native). - ``perception_tools``: the ``exec_*`` perception/vision tool method group. - ``universal_runner``: the LangChain-based ``_run_universal`` loop. - ``run_setup``: flash mode and the pre-loop setup phases of ``run``. - ``native_runner``: the native Gemini reasoning loop of ``run``. - -All historical module-level names (``settings``, ``StorageManager``, -``genai``, ``_run_object_detection``, ``draw_dots``, ``perform_ocr``, -``is_ocr_configured``, ``get_llm``, ``UNIVERSAL_EXPLORER_TOOLS``, ...) remain -importable and patchable here; split-out code resolves them through this -module at call time (see ``artemis.agents.explorer._facade``). """ -import asyncio import json import os from typing import Any, Literal @@ -39,6 +31,7 @@ from typing import Any, Literal from google import genai from google.genai import types +from artemis.agents.explorer.constants import DEFAULT_EXPLORER_MODEL from artemis.agents.explorer.geometry import ( FALLBACK_SCREEN_SIZE, is_valid_norm_point, @@ -50,51 +43,18 @@ from artemis.agents.explorer.perception_tools import PerceptionToolsMixin from artemis.agents.explorer.run_setup import RunSetupMixin from artemis.agents.explorer.screen_index import normalize_text, ScreenElement, ScreenIndex from artemis.agents.explorer.tiers import ExplorerTier, get_tier -from artemis.agents.explorer.tool_declarations import ( - NATIVE_EXPLORER_TOOL_DECLARATIONS, - UNIVERSAL_EXPLORER_TOOLS, -) +from artemis.agents.explorer.tool_declarations import NATIVE_EXPLORER_TOOL_DECLARATIONS from artemis.agents.explorer.universal_runner import UniversalRunnerMixin -from artemis.agents.object_detector.object_detector import _run_object_detection from artemis.config import settings from artemis.context import ArtemisContext -from artemis.data_engine.storage import StorageManager from artemis.data_engine.trace import trace from artemis.graph.state import State -from artemis.llm.reliability import ( - LLMExhaustedError, - LLMPermanentError, - classify_failure, - retry_policy_for, -) -from artemis.services.llm import _record_llm_event, _record_llm_retry, get_llm from artemis.utils.logger import get_logger -from artemis.utils.ocr_api import is_ocr_configured, perform_ocr -from artemis.utils.visualization import draw_dots +from artemis.utils.ocr_api import is_ocr_configured logger = get_logger(__name__) -#: Public surface plus the late-bound collaborators the split-out modules -#: resolve through this facade (``mock.patch`` targets for the test suite). -__all__ = [ - "DEFAULT_EXPLORER_MODEL", - "Explorer", - "NATIVE_EXPLORER_TOOL_DECLARATIONS", - "StorageManager", - "UNIVERSAL_EXPLORER_TOOLS", - "_generate_content_with_reliability", - "_run_object_detection", - "draw_dots", - "genai", - "get_llm", - "is_ocr_configured", - "logger", - "perform_ocr", - "settings", -] - -#: Model used when the LLM configuration carries no Explorer entry. -DEFAULT_EXPLORER_MODEL = "gemini-3.8-flash" +__all__ = ["Explorer"] #: Environment override for Gemini explicit context caching ("true" / "false"). CACHING_ENV_VAR = "ARTEMIS_EXPLORER_CACHING" @@ -104,80 +64,6 @@ CACHING_ENV_VAR = "ARTEMIS_EXPLORER_CACHING" LABEL_TOLERANCE_RATIO = 0.02 -async def _generate_content_with_reliability(operation, *, label: str = "Explorer model call"): - """Run one native google-genai model call under the shared reliability layer. - - Retry decisions are owned by ``classify_failure``/``retry_policy_for`` - (artemis.llm.reliability) instead of a blanket exponential-backoff loop: - non-retryable categories (auth, bad request) are raised immediately as - ``LLMPermanentError`` and exhausted retryable categories surface as - ``LLMExhaustedError``. Only the retry/error-classification policy is - centralized here; the google-genai transport itself stays untouched. - ``operation`` must be a zero-argument callable returning a fresh awaitable - per attempt. - """ - attempt = 0 - while True: - try: - return await operation() - except Exception as call_err: - attempt += 1 - failure = classify_failure(call_err) - if not failure.retryable: - logger.error(f"{label} permanently failed [{failure.category.value}]: {call_err}") - _record_llm_event( - "llm_gave_up", - { - "source": "explorer", - "error": str(call_err)[:1000], - "category": failure.category.value, - "retryable": False, - }, - status="failed", - ) - raise LLMPermanentError( - f"{label} failed [{failure.category.value}]: {call_err}", - failure=failure, - cause=call_err, - ) from call_err - policy = retry_policy_for(failure.category) - if attempt >= policy.max_attempts: - logger.error( - f"{label} exhausted {attempt} attempt(s) [{failure.category.value}]: {call_err}" - ) - _record_llm_event( - "llm_gave_up", - { - "source": "explorer", - "error": str(call_err)[:1000], - "category": failure.category.value, - "retryable": True, - "attempts": attempt, - }, - status="failed", - ) - raise LLMExhaustedError( - f"{label} exhausted {attempt} attempt(s)" - f" [{failure.category.value}]: {call_err}", - failure=failure, - cause=call_err, - ) from call_err - delay = policy.delay_for(attempt) - logger.warning( - f"{label} failed [{failure.category.value}] on attempt" - f" {attempt}/{policy.max_attempts}: {call_err}." - f" Retrying in {delay:.2f}s..." - ) - _record_llm_retry( - str(call_err), - delay, - attempt=attempt, - max_retries=policy.max_attempts, - source="explorer", - ) - await asyncio.sleep(delay) - - class Explorer(PerceptionToolsMixin, UniversalRunnerMixin, RunSetupMixin, NativeRunnerMixin): """Locates UI elements on a screenshot for the ``ask_explorer`` tool. diff --git a/artemis/agents/explorer/native_runner.py b/artemis/agents/explorer/native_runner.py index 4b97f11..3f62c99 100644 --- a/artemis/agents/explorer/native_runner.py +++ b/artemis/agents/explorer/native_runner.py @@ -19,9 +19,9 @@ loop and its named phases (model invocation, submit_answer interception and validation, tool dispatching, failure mapping, resource cleanup), packaged as a mixin consumed by ``Explorer``. Every google-genai call goes through the SDK's ``client.aio`` surface so uploads, cache management and generation -never block the event loop. Patched collaborators (``logger``, -``_generate_content_with_reliability``) are resolved through the facade -module at call time; see ``artemis.agents.explorer._facade``. +never block the event loop. This module also owns +:func:`_generate_content_with_reliability`, the shared reliability wrapper +for native google-genai model calls. """ import asyncio @@ -33,16 +33,94 @@ from typing import TYPE_CHECKING, Any from google.genai import types -from artemis.agents.explorer._facade import facade from artemis.agents.explorer.geometry import is_valid_norm_point from artemis.agents.explorer.tiers import SUBMIT_TOOL from artemis.constants import SAFETY_SETTINGS_BLOCK_NONE from artemis.data_engine.trace import TraceSpan +from artemis.llm.reliability import ( + LLMExhaustedError, + LLMPermanentError, + classify_failure, + retry_policy_for, +) +from artemis.services.llm import _record_llm_event, _record_llm_retry +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) #: Loads a local image as a ``Part`` (inline bytes or a File API reference). ImagePartGetter = Callable[[str], Awaitable[types.Part]] +async def _generate_content_with_reliability(operation, *, label: str = "Explorer model call"): + """Run one native google-genai model call under the shared reliability layer. + + ``operation`` must return a fresh awaitable on each call. Non-retryable + failures raise ``LLMPermanentError``; exhausted retries raise + ``LLMExhaustedError`` according to the shared category policy. + """ + attempt = 0 + while True: + try: + return await operation() + except Exception as call_err: + attempt += 1 + failure = classify_failure(call_err) + if not failure.retryable: + logger.error(f"{label} permanently failed [{failure.category.value}]: {call_err}") + _record_llm_event( + "llm_gave_up", + { + "source": "explorer", + "error": str(call_err)[:1000], + "category": failure.category.value, + "retryable": False, + }, + status="failed", + ) + raise LLMPermanentError( + f"{label} failed [{failure.category.value}]: {call_err}", + failure=failure, + cause=call_err, + ) from call_err + policy = retry_policy_for(failure.category) + if attempt >= policy.max_attempts: + logger.error( + f"{label} exhausted {attempt} attempt(s) [{failure.category.value}]: {call_err}" + ) + _record_llm_event( + "llm_gave_up", + { + "source": "explorer", + "error": str(call_err)[:1000], + "category": failure.category.value, + "retryable": True, + "attempts": attempt, + }, + status="failed", + ) + raise LLMExhaustedError( + f"{label} exhausted {attempt} attempt(s)" + f" [{failure.category.value}]: {call_err}", + failure=failure, + cause=call_err, + ) from call_err + delay = policy.delay_for(attempt) + logger.warning( + f"{label} failed [{failure.category.value}] on attempt" + f" {attempt}/{policy.max_attempts}: {call_err}." + f" Retrying in {delay:.2f}s..." + ) + _record_llm_retry( + str(call_err), + delay, + attempt=attempt, + max_retries=policy.max_attempts, + source="explorer", + ) + await asyncio.sleep(delay) + + def _bare_tool_name(name: str) -> str: """Strips a namespace prefix (``ns:tool``) some models add to tool names.""" return name.split(":")[-1] if ":" in name else name @@ -145,9 +223,8 @@ Initial marked UI elements list (corresponding to numbers ①, ②... in the ima self, client, model_name: str, prompt_template: str, contents: list ): """Creates the explicit cache resource; returns None on failure.""" - _ex = facade() try: - _ex.logger.info("Creating cache resource for Explorer...") + logger.info("Creating cache resource for Explorer...") cached_content = await client.aio.caches.create( model=model_name, config=types.CreateCachedContentConfig( @@ -163,10 +240,10 @@ Initial marked UI elements list (corresponding to numbers ①, ②... in the ima ttl="300s", ), ) - _ex.logger.info(f"Cache resource created successfully: {cached_content.name}") + logger.info(f"Cache resource created successfully: {cached_content.name}") return cached_content except Exception as cache_err: - _ex.logger.error(f"Failed to create cache resource: {cache_err}") + logger.error(f"Failed to create cache resource: {cache_err}") return None async def _run_native( @@ -234,7 +311,6 @@ Initial marked UI elements list (corresponding to numbers ①, ②... in the ima max_turns: int, ) -> str: """Main native reasoning loop: one model turn per iteration.""" - _ex = facade() turn = 0 agent_outcome = "" self.turn_latencies = [] @@ -243,7 +319,7 @@ Initial marked UI elements list (corresponding to numbers ①, ②... in the ima while turn < max_turns: turn += 1 - _ex.logger.info(f"Turn {turn}/{max_turns}: Invoking Native Gemini SDK for Explorer...") + logger.info(f"Turn {turn}/{max_turns}: Invoking Native Gemini SDK for Explorer...") self._prune_historical_images(contents, keep_last=1) @@ -413,7 +489,6 @@ Initial marked UI elements list (corresponding to numbers ①, ②... in the ima thinking_level, ): """Invokes the model for one turn under a trace span; returns response and texts.""" - _ex = facade() ctx = self.ctx start_turn = time.perf_counter() with TraceSpan(name="gemini_explorer_call", ctx=ctx) as span: @@ -425,7 +500,7 @@ Initial marked UI elements list (corresponding to numbers ①, ②... in the ima thinking_level=thinking_level, ) - response = await _ex._generate_content_with_reliability( + response = await _generate_content_with_reliability( lambda: asyncio.wait_for( client.aio.models.generate_content( model=model_name, @@ -476,8 +551,7 @@ Initial marked UI elements list (corresponding to numbers ①, ②... in the ima recorded as a synthetic model text part so the conversation history stays well-formed for the retry. """ - _ex = facade() - _ex.logger.warning(f"Explorer turn {turn}: Model hallucinated plain text. Forcing retry.") + logger.warning(f"Explorer turn {turn}: Model hallucinated plain text. Forcing retry.") model_text = "\n".join(text_parts) if text_parts else "" turn_record["tool_calls"].append( { @@ -573,7 +647,6 @@ Initial marked UI elements list (corresponding to numbers ①, ②... in the ima On validation failure the error is fed back as a tool response and None is returned so the ReAct loop can continue and self-correct. """ - _ex = facade() submit_call = next( (fc for fc in function_calls if _bare_tool_name(fc.name) == SUBMIT_TOOL), None, @@ -586,7 +659,7 @@ Initial marked UI elements list (corresponding to numbers ①, ②... in the ima errors = self._validate_submit_args(function_calls, args, candidates) if errors: - _ex.logger.warning(f"Explorer submit_answer validation failed: {errors}") + logger.warning(f"Explorer submit_answer validation failed: {errors}") turn_record["tool_calls"].append( { "name": SUBMIT_TOOL, @@ -611,7 +684,7 @@ Initial marked UI elements list (corresponding to numbers ①, ②... in the ima # Allow the ReAct loop to continue so the model can self-correct return None - _ex.logger.info("Explorer submit_answer validation passed successfully!") + logger.info("Explorer submit_answer validation passed successfully!") # Candidates submitted by a registered label inherit that element's # bounds; the rest of the model's arguments are passed through as is. outcome_args = dict(args) @@ -637,7 +710,6 @@ Initial marked UI elements list (corresponding to numbers ①, ②... in the ima Image loading failures are logged and skipped: the textual result is still delivered, so one unreadable annotation cannot sink the turn. """ - _ex = facade() tool_response_parts.append( types.Part.from_function_response(name=name, response={result_key: res.get("text")}) ) @@ -646,7 +718,7 @@ Initial marked UI elements list (corresponding to numbers ①, ②... in the ima try: tool_response_parts.append(await get_image_part(img_p)) except Exception as e: - _ex.logger.warning(f"Failed to load image response part for {img_p}: {e}") + logger.warning(f"Failed to load image response part for {img_p}: {e}") async def _execute_native_tool( self, @@ -730,7 +802,6 @@ Initial marked UI elements list (corresponding to numbers ①, ②... in the ima absent from the declarations, because models occasionally call tools they were never offered. """ - _ex = facade() for fc in function_calls: name = _bare_tool_name(fc.name) if name == SUBMIT_TOOL: @@ -738,7 +809,7 @@ Initial marked UI elements list (corresponding to numbers ①, ②... in the ima args = fc.args or {} tool_call_trace = {"name": name, "args": args} if name in self.denylisted_tools: - _ex.logger.warning( + logger.warning( f"Explorer attempted to call denylisted tool '{name}'. Blocking execution." ) error = f"Tool '{name}' is denylisted and unavailable." @@ -749,14 +820,14 @@ Initial marked UI elements list (corresponding to numbers ①, ②... in the ima ) continue - _ex.logger.info(f"Explorer executing tool '{name}' sequentially...") + logger.info(f"Explorer executing tool '{name}' sequentially...") try: await self._execute_native_tool( name, args, get_image_part, tool_call_trace, tool_response_parts ) except Exception as e: - _ex.logger.error(f"Explorer tool {name} execution failed: {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)}) @@ -766,8 +837,7 @@ Initial marked UI elements list (corresponding to numbers ①, ②... in the ima def _format_native_failure(self, e: Exception, model_name: str) -> str: """Maps a loop failure to the user-facing fallback outcome JSON.""" - _ex = facade() - _ex.logger.error(f"Explorer execution loop failed: {e}") + logger.error(f"Explorer execution loop failed: {e}") err_msg = str(e) if "preempted" in err_msg.lower(): @@ -794,16 +864,15 @@ Initial marked UI elements list (corresponding to numbers ①, ②... in the ima async def _cleanup_native_resources(self, client, cached_content, uploaded_files: list) -> None: """Deletes the cache resource and File API uploads created for this run.""" - _ex = facade() if cached_content: try: - _ex.logger.info(f"Deleting cache resource: {cached_content.name}") + logger.info(f"Deleting cache resource: {cached_content.name}") await client.aio.caches.delete(name=cached_content.name) except Exception as cleanup_cache_err: - _ex.logger.warning(f"Failed to delete cache resource: {cleanup_cache_err}") + logger.warning(f"Failed to delete cache resource: {cleanup_cache_err}") for file_ref in uploaded_files: try: await client.aio.files.delete(name=file_ref.name) except Exception as cleanup_err: - _ex.logger.warning(f"Failed to delete uploaded file {file_ref.name}: {cleanup_err}") - _ex.logger.info(f"Explorer turn latencies: {self.turn_latencies}") + logger.warning(f"Failed to delete uploaded file {file_ref.name}: {cleanup_err}") + logger.info(f"Explorer turn latencies: {self.turn_latencies}") diff --git a/artemis/agents/explorer/perception_tools.py b/artemis/agents/explorer/perception_tools.py index 86ae6a0..a8a6350 100644 --- a/artemis/agents/explorer/perception_tools.py +++ b/artemis/agents/explorer/perception_tools.py @@ -19,9 +19,7 @@ group plus its search helpers, packaged as a mixin consumed by ``Explorer``. Text search, the coordinate audit and the OCR inventory all answer from the run's in-memory :class:`~artemis.agents.explorer.screen_index.ScreenIndex`; no Data Engine round trip is involved, so an unsynced frame never degrades -to "no UI data". Patched collaborators (``settings``, ``draw_dots``, -``_run_object_detection``, ``logger``) are resolved through the facade module -at call time; see ``artemis.agents.explorer._facade``. +to "no UI data". """ import asyncio @@ -33,11 +31,16 @@ from typing import TYPE_CHECKING, Any import cv2 -from artemis.agents.explorer._facade import facade from artemis.agents.explorer.geometry import norm_to_pixel, pixel_to_norm from artemis.agents.explorer.screen_index import ScreenElement, ScreenIndex, normalize_text 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.data_engine.trace import trace +from artemis.utils.logger import get_logger +from artemis.utils.visualization import draw_dots + +logger = get_logger(__name__) #: Text-search answer when the screen could not be indexed at all. NO_UI_TREE_MESSAGE = "No UI-tree data is available for this screen; rely on visual detection." @@ -49,7 +52,6 @@ def load_detector_templates() -> tuple[list[str], float]: Shared by ``exec_detect_objects`` and the flash-mode ``run`` path; the two call sites historically carried an identical inline copy of this block. """ - _ex = facade() detector_prompt_path = Path(__file__).parent.parent / "object_detector" / "object_detector.json" global_timeout = 30.0 try: @@ -61,7 +63,7 @@ def load_detector_templates() -> tuple[list[str], float]: if isinstance(raw_timeout, (int, float)) and raw_timeout > 0: global_timeout = float(raw_timeout) except Exception as e: - _ex.logger.warning(f"Failed to load detector prompt config: {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 = "" @@ -127,8 +129,7 @@ class PerceptionToolsMixin: def _next_output_path(self, subdir: str) -> Path: """Next free ``images//_.jpg`` under the traces directory.""" - _ex = facade() - out_dir = Path(_ex.settings.TRACES_PATH) / "images" / subdir + out_dir = Path(settings.TRACES_PATH) / "images" / subdir out_dir.mkdir(parents=True, exist_ok=True) stem = self.image_name or "temp_image" max_seq = 0 @@ -147,15 +148,14 @@ class PerceptionToolsMixin: textual result is still delivered, so one bad annotation never sinks the tool call. """ - _ex = facade() if not self.screenshot_path: return None try: output_path = self._next_output_path(subdir) - _ex.draw_dots(self.screenshot_path, points, labels, str(output_path), color=color) + draw_dots(self.screenshot_path, points, labels, str(output_path), color=color) return str(output_path) except Exception as e: - _ex.logger.warning(f"Failed to draw {subdir} dots: {e}") + logger.warning(f"Failed to draw {subdir} dots: {e}") return None # ------------------------------------------------------------------ # @@ -237,7 +237,6 @@ class PerceptionToolsMixin: prefix: str = "D", color: str = "red", ) -> dict: - _ex = facade() try: if not self.screenshot_path: return { @@ -257,7 +256,7 @@ class PerceptionToolsMixin: # Load object detector templates templates, global_timeout = load_detector_templates() - result = await _ex._run_object_detection( + result = await _run_object_detection( self.ctx, target_path, queries, @@ -280,7 +279,7 @@ class PerceptionToolsMixin: "image_path": str(target_path), } - base_dir = _ex.settings.TRACES_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) @@ -341,7 +340,7 @@ class PerceptionToolsMixin: f'[{label_id}] "{item.get("label")}" at [{x_orig_norm},{y_orig_norm}]' ) - _ex.draw_dots( + draw_dots( self.screenshot_path, points, d_labels, @@ -361,7 +360,7 @@ class PerceptionToolsMixin: } except Exception as e: - _ex.logger.warning(f"Failed to process detection output for annotation: {e}") + logger.warning(f"Failed to process detection output for annotation: {e}") return {"text": str(result), "image_path": None} except Exception as e: return { @@ -569,7 +568,6 @@ class PerceptionToolsMixin: y_max: int, zoom_factor: float = 2.0, ) -> dict: - _ex = facade() try: if not self.screenshot_path: return { @@ -613,7 +611,7 @@ class PerceptionToolsMixin: if new_w > 0 and new_h > 0: cropped = cv2.resize(cropped, (new_w, new_h)) - base_dir = _ex.settings.TRACES_PATH + 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) diff --git a/artemis/agents/explorer/run_setup.py b/artemis/agents/explorer/run_setup.py index 3933b88..60f7f55 100644 --- a/artemis/agents/explorer/run_setup.py +++ b/artemis/agents/explorer/run_setup.py @@ -19,11 +19,7 @@ detection path, the per-run screen loading shared by every tier (image record resolution, screen dimensions, fused hierarchy, in-memory screen index) and the named preparation phases executed before a reasoning loop (image pool, model parameters, tier-aware prompt construction, initial screenshot -annotation), packaged as a mixin consumed by ``Explorer``. Patched -collaborators (``settings``, ``StorageManager``, ``is_ocr_configured``, -``perform_ocr``, ``draw_dots``, ``_run_object_detection``, ``logger``) are -resolved through the facade module at call time; see -``artemis.agents.explorer._facade``. +annotation), packaged as a mixin consumed by ``Explorer``. """ import base64 @@ -37,15 +33,24 @@ from typing import TYPE_CHECKING, Any import httpx -from artemis.agents.explorer._facade import facade -from artemis.agents.explorer.constants import EXECUTION_CONSTRAINT_TEMPLATE +from artemis.agents.explorer.constants import ( + DEFAULT_EXPLORER_MODEL, + EXECUTION_CONSTRAINT_TEMPLATE, +) from artemis.agents.explorer.geometry import resolve_screen_size from artemis.agents.explorer.perception_tools import load_detector_templates from artemis.agents.explorer.screen_index import ScreenElement, ScreenIndex from artemis.agents.explorer.tiers import SUBMIT_TOOL, ExplorerTier +from artemis.agents.object_detector.object_detector import _run_object_detection +from artemis.config import settings +from artemis.data_engine.storage import StorageManager from artemis.graph.state import State +from artemis.utils.logger import get_logger +from artemis.utils.ocr_api import is_ocr_configured, perform_ocr from artemis.utils.ocr_xml_fusion import fuse_ocr_with_xml -from artemis.utils.visualization import format_minimal_list_with_elements +from artemis.utils.visualization import draw_dots, format_minimal_list_with_elements + +logger = get_logger(__name__) #: Prompt bullets may be plain strings or ``{"text": ..., "requires": [...]}`` #: objects; the latter are emitted only when every listed tool is exposed. @@ -109,13 +114,12 @@ class RunSetupMixin: Several targets may be listed in ``query`` separated by ``|``; each detection becomes a ``D`` candidate described by its detected label. """ - _ex = facade() templates, global_timeout = load_detector_templates() queries = [q.strip() for q in query.split("|") if q.strip()] try: - result = await _ex._run_object_detection( + result = await _run_object_detection( self.ctx, screenshot_path, queries, @@ -140,12 +144,11 @@ class RunSetupMixin: ensure_ascii=False, ) except Exception as e: - _ex.logger.error(f"Flash mode object detection failed: {e}") + logger.error(f"Flash mode object detection failed: {e}") return self._failure_outcome(f"Flash mode detection error: {e}") def _resolve_image_record(self, screenshot_path: str) -> tuple: """Computes the screenshot hash and looks it up in the Data Engine DB.""" - _ex = facade() image_name = None record = None try: @@ -154,24 +157,24 @@ class RunSetupMixin: for byte_block in iter(lambda: f.read(4096), b""): sha256_hash.update(byte_block) computed_hash = sha256_hash.hexdigest() - _ex.logger.info(f"Computed screenshot hash: {computed_hash}") + logger.info(f"Computed screenshot hash: {computed_hash}") # Check if it exists in DB - db_path = _ex.settings.DATA_ENGINE_DB_PATH - base_dir = _ex.settings.TRACES_PATH - storage = _ex.StorageManager(db_path, base_dir) + 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: - _ex.logger.warning( + logger.warning( f"Image hash {computed_hash} not found in Data Engine DB." " Data Engine might not be synced yet." ) except Exception as e: - _ex.logger.warning(f"Failed to compute hash or check DB: {e}") + logger.warning(f"Failed to compute hash or check DB: {e}") return image_name, record @@ -216,7 +219,6 @@ class RunSetupMixin: def _resolve_model_params(self) -> tuple: """Resolves model name, temperature, thinking level, and fallback model.""" - _ex = facade() llm_config = getattr(self.ctx, "llm_config", None) llm_cfg = getattr(llm_config, "explorer", None) if llm_config else None temperature = 0.1 @@ -235,7 +237,7 @@ class RunSetupMixin: if "/" in fallback_model: fallback_model = fallback_model.split("/")[-1] else: - model_name = _ex.DEFAULT_EXPLORER_MODEL + model_name = DEFAULT_EXPLORER_MODEL return model_name, temperature, thinking_level, fallback_model @@ -303,32 +305,27 @@ class RunSetupMixin: HTTP client backing on-the-fly OCR is created here on first need and released by ``Explorer.run`` regardless of the engine used. """ - _ex = facade() fused_xml: Any = [] if record and getattr(record, "ui_tree", None): ui_tree = record.ui_tree ocr_results = getattr(record, "ocr_result", None) if ocr_results is None: - if allow_ocr and _ex.is_ocr_configured(): + if allow_ocr and is_ocr_configured(): try: - _ex.logger.info( - "Previous screenshot OCR is missing. Running OCR on-the-fly..." - ) + logger.info("Previous screenshot OCR is missing. Running OCR on-the-fly...") with open(screenshot_path, "rb") as img_file: img_b64 = base64.b64encode(img_file.read()).decode("utf-8") if self.http_client is None: self.http_client = httpx.AsyncClient() - ocr_results = await _ex.perform_ocr(img_b64, client=self.http_client) + ocr_results = await perform_ocr(img_b64, client=self.http_client) except Exception as ocr_err: - _ex.logger.error( - f"On-the-fly OCR failed for previous screenshot: {ocr_err}" - ) + logger.error(f"On-the-fly OCR failed for previous screenshot: {ocr_err}") ocr_results = [] else: ocr_results = [] fused_xml = fuse_ocr_with_xml(ui_tree, ocr_results or []) - _ex.logger.info("Successfully loaded and fused UI hierarchy for previous screenshot.") + logger.info("Successfully loaded and fused UI hierarchy for previous screenshot.") elif screenshot_path == getattr(state, "latest_screenshot", None): fused_xml = getattr(state, "latest_ui_hierarchy", None) return list(fused_xml) if isinstance(fused_xml, (list, tuple)) else [] @@ -341,13 +338,10 @@ class RunSetupMixin: Every numbered label is registered against its element so a candidate the model later submits by number inherits real bounds. """ - _ex = facade() ctx = self.ctx marked_path = None try: - _ex.logger.info( - "Explorer self-annotating initial screenshot using latest_ui_hierarchy..." - ) + logger.info("Explorer self-annotating initial screenshot using latest_ui_hierarchy...") formatted_list, elements, labels = format_minimal_list_with_elements( fused_xml, self.width, self.height ) @@ -363,7 +357,7 @@ class RunSetupMixin: else None ) if not base_dir: - base_dir = _ex.settings.TRACES_PATH + base_dir = settings.TRACES_PATH images_dir = base_dir / "images" initial_marked_dir = images_dir / "initial_marked" initial_marked_dir.mkdir(parents=True, exist_ok=True) @@ -380,12 +374,12 @@ class RunSetupMixin: marked_path = initial_marked_dir / f"{image_name or 'temp_image'}_{seq}.jpg" # Draw dots on the raw screenshot - _ex.draw_dots(screenshot_path, points, labels, str(marked_path)) - _ex.logger.info( + draw_dots(screenshot_path, points, labels, str(marked_path)) + logger.info( f"Successfully drew {len(points)} dots and saved marked image to {marked_path}" ) except Exception as e: - _ex.logger.error(f"Failed to self-annotate initial screenshot: {e}") + logger.error(f"Failed to self-annotate initial screenshot: {e}") return minimal_list, marked_path def _prepare_initial_annotation( diff --git a/artemis/agents/explorer/tool_declarations.py b/artemis/agents/explorer/tool_declarations.py index e8c5969..16141df 100644 --- a/artemis/agents/explorer/tool_declarations.py +++ b/artemis/agents/explorer/tool_declarations.py @@ -17,8 +17,7 @@ Split out of ``artemis.agents.explorer.explorer`` as pure data: ``UNIVERSAL_EXPLORER_TOOLS`` (OpenAI-style dict schemas for the universal LangChain path) and ``NATIVE_EXPLORER_TOOL_DECLARATIONS`` (google-genai -``FunctionDeclaration`` objects for the native Gemini path). Both are -re-exported by the ``explorer`` facade module. +``FunctionDeclaration`` objects for the native Gemini path). """ from google.genai import types diff --git a/artemis/agents/explorer/universal_runner.py b/artemis/agents/explorer/universal_runner.py index 6dbfc4c..9a123ca 100644 --- a/artemis/agents/explorer/universal_runner.py +++ b/artemis/agents/explorer/universal_runner.py @@ -15,10 +15,7 @@ """Universal (LangChain ChatModel) execution loop for the Explorer agent. Split out of ``artemis.agents.explorer.explorer``: the ``_run_universal`` -reasoning loop, packaged as a mixin consumed by ``Explorer``. Patched -collaborators (``get_llm``, ``logger``, ``UNIVERSAL_EXPLORER_TOOLS``) are -resolved through the facade module at call time; see -``artemis.agents.explorer._facade``. +reasoning loop, packaged as a mixin consumed by ``Explorer``. """ import asyncio @@ -33,10 +30,14 @@ from langchain_core.messages import ( ToolMessage, ) -from artemis.agents.explorer._facade import facade from artemis.agents.explorer.geometry import is_valid_norm_point from artemis.agents.explorer.tiers import SUBMIT_TOOL +from artemis.agents.explorer.tool_declarations import UNIVERSAL_EXPLORER_TOOLS from artemis.graph.state import State +from artemis.services.llm import get_llm +from artemis.utils.logger import get_logger + +logger = get_logger(__name__) class UniversalRunnerMixin: @@ -72,9 +73,8 @@ class UniversalRunnerMixin: def _universal_exposed_tools(self) -> list[dict[str, Any]]: """OpenAI-style tool schemas the model may see (tier-hidden ones excluded).""" - _ex = facade() hidden = self._hidden_tool_names() - return [t for t in _ex.UNIVERSAL_EXPLORER_TOOLS if t["function"]["name"] not in hidden] + return [t for t in UNIVERSAL_EXPLORER_TOOLS if t["function"]["name"] not in hidden] async def _run_universal( self, @@ -88,8 +88,7 @@ class UniversalRunnerMixin: max_turns: int, ) -> str: """Executes the Explorer reasoning loop via a universal LangChain ChatModel.""" - _ex = facade() - llm = _ex.get_llm(self.ctx, name="explorer") + llm = get_llm(self.ctx, name="explorer") exposed_tools = self._universal_exposed_tools() bound_llm = llm.bind_tools(exposed_tools) @@ -124,7 +123,7 @@ class UniversalRunnerMixin: tool_calls = getattr(response, "tool_calls", None) or [] if not tool_calls: - _ex.logger.warning(f"Universal Explorer turn {turn}: No tool calls generated.") + logger.warning(f"Universal Explorer turn {turn}: No tool calls generated.") if is_final_turn: return json.dumps( { @@ -256,7 +255,6 @@ class UniversalRunnerMixin: the images ride on a follow-up human turn instead; unreadable files are skipped so a missing annotation never aborts the loop. """ - _ex = facade() blocks: list[dict[str, Any]] = [] tool_names: list[str] = [] for tool_name, path in images: @@ -264,7 +262,7 @@ class UniversalRunnerMixin: with open(path, "rb") as f: img_b64 = base64.b64encode(f.read()).decode("utf-8") except OSError as read_err: - _ex.logger.warning(f"Skipping unreadable tool image {path}: {read_err}") + logger.warning(f"Skipping unreadable tool image {path}: {read_err}") continue if tool_name not in tool_names: tool_names.append(tool_name) @@ -283,7 +281,6 @@ class UniversalRunnerMixin: self, tool_calls: list, messages: list[BaseMessage], turn: int ) -> None: """Executes non-submit tool calls, appending ToolMessages and their images.""" - _ex = facade() images: list[tuple[str, str]] = [] for tc in tool_calls: name = tc.get("name") @@ -306,7 +303,7 @@ class UniversalRunnerMixin: messages.append(ToolMessage(tool_call_id=call_id, name=name, content=res_text)) images.extend((name, path) for path in self._tool_image_paths(res)) except Exception as tool_err: - _ex.logger.error(f"Error executing tool '{name}': {tool_err}") + logger.error(f"Error executing tool '{name}': {tool_err}") messages.append( ToolMessage( tool_call_id=call_id, diff --git a/artemis/agents/operator/operator.py b/artemis/agents/operator/operator.py index dcdfeb4..359fc87 100644 --- a/artemis/agents/operator/operator.py +++ b/artemis/agents/operator/operator.py @@ -82,7 +82,7 @@ from artemis.agents.operator.prompts import ( InjectedInstructionPromptComponent, render_transcript_static_system, ) -from artemis.agents.operator.prompt_builder import load_operator_prompts +from artemis.agents.operator.prompts import load_operator_prompts class OperatorNode: @@ -101,8 +101,7 @@ class OperatorNode: try: self.prompts = load_operator_prompts() except (OSError, ValueError, json.JSONDecodeError) as e: - logger.error(f"Failed to load operator prompts: {e}") - self.prompts = {} + raise RuntimeError(f"Failed to load operator prompts: {e}") from e # M2 transcript flag (agent.memory.transcript.enabled). Off keeps the # legacy 2-message prompt path byte-for-byte; an explicit @@ -488,8 +487,7 @@ class OperatorNode: 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 type(additional_kwargs) is dict and additional_kwargs.get("function_call") ): fc = additional_kwargs["function_call"] diff --git a/artemis/agents/operator/prompt_builder.py b/artemis/agents/operator/prompt_builder.py deleted file mode 100644 index 5a09ce5..0000000 --- a/artemis/agents/operator/prompt_builder.py +++ /dev/null @@ -1,71 +0,0 @@ -# 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. - -"""Compatibility prompt builder backed by the canonical Operator templates.""" - -import json -from functools import lru_cache -from pathlib import Path - -from jinja2 import Template - -OPERATOR_HUMAN_TEMPLATE = """Goal: {{ goal }} -Active Sub-Goal: {{ sub_goal }} -Current Turn: {{ current_turn }} -Execution History: -{{ history }} -""" - - -@lru_cache(maxsize=1) -def load_operator_prompts() -> dict[str, str]: - """Load the single canonical prompt source shared by all Operator entry points.""" - prompts_path = Path(__file__).with_name("operator.json") - return json.loads(prompts_path.read_text(encoding="utf-8")) - - -class OperatorPromptBuilder: - """Builds prompt messages and multimodal contents for the Operator Agent.""" - - @classmethod - def build_system_message( - cls, - template_name: str = "main_template", - available_tools: frozenset[str] | None = None, - ) -> str: - from artemis.agents.operator.prompts import apply_operator_prompt_contract - - prompts = load_operator_prompts() - try: - return apply_operator_prompt_contract( - prompts[template_name], available_tools=available_tools - ) - except KeyError as exc: - raise ValueError(f"Unknown Operator template: {template_name}") from exc - - @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 index 28b4fa2..13aa0af 100644 --- a/artemis/agents/operator/prompts.py +++ b/artemis/agents/operator/prompts.py @@ -13,7 +13,9 @@ # limitations under the License. import base64 +from functools import lru_cache import io +import json from pathlib import Path from jinja2 import Environment, StrictUndefined, Template @@ -36,6 +38,14 @@ logger = get_logger(__name__) from artemis.agents.prompt_assembly import render_tool_enum, resolve_available from artemis.mcp.action_specs import OPERATOR_SHELL_ORDER + +@lru_cache(maxsize=1) +def load_operator_prompts() -> dict[str, str]: + """Load and cache the Operator templates from operator.json.""" + prompts_path = Path(__file__).with_name("operator.json") + return json.loads(prompts_path.read_text(encoding="utf-8")) + + # --- Assembled tool references ------------------------------------------------------- # The operator.json templates carry availability slots rendered by # ``apply_operator_prompt_contract``: ``[[ tool_enum(...) ]]`` expands an ordered, diff --git a/artemis/clients/ui_automator_client.py b/artemis/clients/ui_automator_client.py index 09b2c7c..b5d3b7a 100644 --- a/artemis/clients/ui_automator_client.py +++ b/artemis/clients/ui_automator_client.py @@ -31,6 +31,7 @@ from PIL import Image from pydantic import BaseModel import uiautomator2 as u2 +from artemis.runtime.adb_endpoint import adb_command from artemis.runtime.awake_service import ensure_device_awake from artemis.utils.logger import get_logger @@ -173,7 +174,7 @@ 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"], + adb_command(["-s", device_id, "shell", "pm", "list", "packages"]), stdin=subprocess.DEVNULL, capture_output=True, text=True, @@ -198,17 +199,18 @@ 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, - ], + adb_command( + [ + "-s", + device_id, + "shell", + "pm", + "uninstall", + "--user", + "0", + pkg, + ] + ), stdin=subprocess.DEVNULL, capture_output=True, text=True, @@ -285,13 +287,24 @@ class UIAutomatorClient: if self._awake_strategy is None: self._awake_strategy = ensure_device_awake(self._device_id) - # Connect to device + # Retry connections while ADB recovers from a reboot or USB disconnect. logger.info(f"Connecting UIAutomator2 to device: {self._device_id}") - try: - self._device = u2.connect(self._device_id) - except Exception: + last_error: Exception | None = None + for attempt in range(3): + if attempt: + time.sleep(1.0 * attempt) + try: + self._device = u2.connect(self._device_id) + break + except Exception as exc: + last_error = exc + logger.warning( + f"UIAutomator2 connect attempt {attempt + 1}/3 to" + f" {self._device_id} failed: {exc}" + ) + else: self._awake_strategy = None - raise + raise last_error logger.info("UIAutomator2 connected successfully") return self._device @@ -380,7 +393,7 @@ class UIAutomatorClient: """ try: result = subprocess.run( - ["adb", "-s", self._device_id, "exec-out", "screencap", "-p"], + adb_command(["-s", self._device_id, "exec-out", "screencap", "-p"]), stdin=subprocess.DEVNULL, capture_output=True, timeout=10, diff --git a/artemis/core/tool_declaration.py b/artemis/core/tool_declaration.py index c40acae..b5e4f80 100644 --- a/artemis/core/tool_declaration.py +++ b/artemis/core/tool_declaration.py @@ -18,7 +18,7 @@ from typing import Any class ToolDeclaration(dict): - """Standardized tool declaration adhering to JSON Schema format with property accessors.""" + """OpenAI function schema with name, description and parameters properties.""" def __init__(self, name: str, description: str, parameters: dict[str, Any]): super().__init__( @@ -28,19 +28,16 @@ class ToolDeclaration(dict): "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", "") + return self.get("function", {}).get("name", "") @property def description(self) -> str: - return self.get("description") or self.get("function", {}).get("description", "") + return self.get("function", {}).get("description", "") @property def parameters(self) -> dict[str, Any]: - return self.get("parameters") or self.get("function", {}).get("parameters", {}) + return self.get("function", {}).get("parameters", {}) diff --git a/artemis/graph/graph.py b/artemis/graph/graph.py index f49396f..b4d023a 100644 --- a/artemis/graph/graph.py +++ b/artemis/graph/graph.py @@ -718,6 +718,50 @@ async def _process_plan_write( return result +async def _intercept_plan_write( + ctx: ArtemisContext, + original_tool, + args: dict, + tool_call_id, + state, + *, + declared_intent: bool, + content_after_fallback: str = "", +): + """Run a note write and process any resulting task-plan changes.""" + rejection = _reject_checker_note_write(args["key"], tool_call_id) + if rejection is not None: + return rejection + if args["key"] != "task_plan" or not ctx.data_engine: + return await invoke_tool_with_injection( + original_tool, args, tool_call_id, state, record_trace=False + ) + + 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, args, tool_call_id, state, record_trace=False + ) + + content_after = content_after_fallback + if task_plan_path.exists(): + content_after = task_plan_path.read_text(encoding="utf-8") + + return await _process_plan_write( + ctx, + state, + task_plan_path, + content_before, + content_after, + result, + declared_intent=declared_intent, + ) + + def wrap_note_tool(ctx: ArtemisContext, original_tool): async def wrapped_note_tool( @@ -727,37 +771,14 @@ def wrap_note_tool(ctx: ArtemisContext, original_tool): state: Annotated[State, InjectedState] = None, ): """Wrapped tool to intercept task plan updates.""" - rejection = _reject_checker_note_write(key, tool_call_id) - if rejection is not None: - return rejection - args = {"key": key, "content": content} - if key != "task_plan" or not ctx.data_engine: - return await invoke_tool_with_injection( - original_tool, args, tool_call_id, state, record_trace=False - ) - - 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, args, 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") - - return await _process_plan_write( + return await _intercept_plan_write( ctx, + original_tool, + {"key": key, "content": content}, + tool_call_id, state, - task_plan_path, - content_before, - content_after, - result, declared_intent=False, + content_after_fallback=content, ) return StructuredTool.from_function( @@ -786,36 +807,12 @@ def wrap_update_note_tool(ctx: ArtemisContext, original_tool): state: Annotated[State, InjectedState] = None, ): """Wrapped tool to intercept task plan updates via update_note.""" - rejection = _reject_checker_note_write(key, tool_call_id) - if rejection is not None: - return rejection - args = {"key": key, "target": target, "replacement": replacement} - if key != "task_plan" or not ctx.data_engine: - return await invoke_tool_with_injection( - original_tool, args, tool_call_id, state, record_trace=False - ) - - 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, args, tool_call_id, state, record_trace=False - ) - - content_after = "" - if task_plan_path.exists(): - content_after = task_plan_path.read_text(encoding="utf-8") - - return await _process_plan_write( + return await _intercept_plan_write( ctx, + original_tool, + {"key": key, "target": target, "replacement": replacement}, + tool_call_id, state, - task_plan_path, - content_before, - content_after, - result, declared_intent=True, ) diff --git a/artemis/llm/router.py b/artemis/llm/router.py index be89df9..69e9c3f 100644 --- a/artemis/llm/router.py +++ b/artemis/llm/router.py @@ -20,6 +20,7 @@ Anthropic, OpenRouter, and local Ollama/vLLM endpoints. """ from enum import StrEnum +import hashlib import os from typing import Any @@ -46,10 +47,19 @@ class ModelProvider(StrEnum): @classmethod def from_string(cls, val: Any) -> "ModelProvider": - """Converts any string, enum, or mock representation into a canonical ModelProvider.""" + """Converts a string or enum into a canonical ModelProvider. + + ``None``/empty means "use the default" (Google); an unrecognized + non-empty value raises instead of silently routing to Gemini with + the wrong credentials. + """ if isinstance(val, cls): return val + if val is None: + return cls.GOOGLE s = str(val).lower().replace("_", "").replace("-", "").strip() + if not s: + return cls.GOOGLE mapping = { "google": cls.GOOGLE, "gemini": cls.GOOGLE, @@ -65,7 +75,12 @@ class ModelProvider(StrEnum): "vllm": cls.VLLM, "custom": cls.CUSTOM, } - return mapping.get(s, cls.GOOGLE) + provider = mapping.get(s) + if provider is None: + raise ValueError( + f"Unknown LLM provider {val!r}. Valid providers: {sorted(set(mapping))}" + ) + return provider class ModelEndpoint(BaseModel): @@ -97,11 +112,15 @@ class ModelEndpoint(BaseModel): ) def cache_key(self) -> tuple: - """Returns a deterministic hashable cache key for connection pooling.""" + """Key the client cache by endpoint settings, hashing the API key.""" + api_key_digest = ( + hashlib.sha256(self.api_key.encode("utf-8")).hexdigest()[:16] if self.api_key else None + ) return ( self.provider.value, self.model_name, self.temperature, + self.max_tokens, self.timeout_seconds, self.thinking_budget, self.thinking_level, @@ -109,6 +128,7 @@ class ModelEndpoint(BaseModel): self.reasoning_effort, self.enable_grounding, self.api_base, + api_key_digest, ) diff --git a/artemis/runtime/adb_endpoint.py b/artemis/runtime/adb_endpoint.py index 1562521..9729f96 100644 --- a/artemis/runtime/adb_endpoint.py +++ b/artemis/runtime/adb_endpoint.py @@ -166,3 +166,8 @@ def current_adb_endpoint() -> AdbEndpoint: host = settings.ADB_HOST or os.environ.get("ADB_HOST") or DEFAULT_ADB_HOST port = settings.ADB_PORT or os.environ.get("ADB_PORT") or DEFAULT_ADB_PORT return AdbEndpoint.create(str(host), int(port)) + + +def adb_command(arguments: Sequence[str]) -> list[str]: + """Resolve the ADB binary and add the configured server's host and port.""" + return AdbSession(current_adb_endpoint()).command(arguments) diff --git a/artemis/runtime/awake_lease.py b/artemis/runtime/awake_lease.py index 75cd625..5f03ab7 100644 --- a/artemis/runtime/awake_lease.py +++ b/artemis/runtime/awake_lease.py @@ -28,6 +28,7 @@ import time import uuid from artemis.config.paths import get_temp_dir +from artemis.runtime.adb_endpoint import adb_command from artemis.runtime.process_probe import pid_is_alive from artemis.utils.logger import get_logger @@ -148,7 +149,7 @@ class ScreenAwakeLease: def _run(self, args: list[str], description: str) -> subprocess.CompletedProcess[str] | None: try: result = subprocess.run( - ["adb", "-s", self.device_id, *args], + adb_command(["-s", self.device_id, *args]), stdin=subprocess.DEVNULL, capture_output=True, text=True, diff --git a/artemis/runtime/awake_service.py b/artemis/runtime/awake_service.py index 12f4e37..2de284e 100644 --- a/artemis/runtime/awake_service.py +++ b/artemis/runtime/awake_service.py @@ -25,6 +25,7 @@ from typing import Iterable from adbutils import AdbClient from artemis.config import settings +from artemis.runtime.adb_endpoint import adb_command from artemis.runtime.awake_lease import ScreenAwakeLease from artemis.utils.logger import get_logger @@ -47,7 +48,7 @@ def _run_awake_adb_command( """Run one non-fatal ADB command used by the USB awake policy.""" try: result = subprocess.run( - ["adb", "-s", device_id, *args], + adb_command(["-s", device_id, *args]), stdin=subprocess.DEVNULL, capture_output=True, text=True, diff --git a/artemis/runtime/device_lock.py b/artemis/runtime/device_lock.py index fcc96f0..9579d23 100644 --- a/artemis/runtime/device_lock.py +++ b/artemis/runtime/device_lock.py @@ -261,6 +261,20 @@ class DeviceExecutionLock: payload["lock_scope"] = str(lock_scope) return payload + @classmethod + def _atomic_write_json(cls, path: Path, payload: dict) -> bool: + """Replace a ticket atomically so queue scans cannot read partial JSON.""" + temp_path = path.with_name(f"{path.name}.{uuid.uuid4().hex}.tmp") + try: + temp_path.write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8") + if not cls._safe_replace(temp_path, path): + return False + except OSError: + return False + finally: + cls._safe_unlink(temp_path) + return True + @classmethod def _ticket_targets_device( cls, @@ -454,11 +468,7 @@ class DeviceExecutionLock: effective_scope = lock_scope or os.getenv(cls.LOCK_SCOPE_ENV) or None if effective_scope: payload["lock_scope"] = str(effective_scope) - try: - path.write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8") - except OSError: - return False - return True + return cls._atomic_write_json(path, payload) def _find_reserved_ticket(self) -> Path | None: if not self.queue_ticket: @@ -482,7 +492,10 @@ class DeviceExecutionLock: ingress=self.ingress, lock_scope=self.lock_scope, ) - path.write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8") + if not self._atomic_write_json(path, payload): + raise DeviceBusyError( + "The Artemis queue reservation could not be claimed (ticket rewrite failed)." + ) return path self.queue_ticket = self.reserve( @@ -567,6 +580,15 @@ class DeviceExecutionLock: return 0 if self.max_concurrency is not None: return self.max_concurrency + return self.resolve_env_concurrency() + + @classmethod + def resolve_env_concurrency(cls) -> int: + """Resolve the concurrency limit from the environment. + + Shared with the admin scheduler. Returns 0 for per-device concurrency, + or a positive global limit. + """ env_mode = os.environ.get("ARTEMIS_CONCURRENCY_MODE", "").strip().lower() if env_mode in ("global", "serial", "1"): return 1 @@ -690,9 +712,14 @@ class DeviceExecutionLock: return try: owner = self._read_owner(self.path) - if owner is not None and ( - owner.token == self.token - or (self.session_id and owner.session_id == self.session_id) + # A matching session id cannot release another process's lease. + if ( + owner is not None + and owner.pid == os.getpid() + and ( + owner.token == self.token + or (self.session_id and owner.session_id == self.session_id) + ) ): self._safe_unlink(self.path) finally: @@ -933,16 +960,7 @@ class DeviceExecutionLock: "session_id": str(session_id), "ingress": ingress or owner.ingress or "sdk", } - temp_path = path.with_name(f"{path.name}.{uuid.uuid4().hex}.tmp") - try: - temp_path.write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8") - if not cls._safe_replace(temp_path, path): - return False - except OSError: - return False - finally: - cls._safe_unlink(temp_path) - return True + return cls._atomic_write_json(path, payload) @classmethod def cleanup_stale_locks(cls, device_id: str | None = None) -> int: diff --git a/artemis/runtime/device_pool.py b/artemis/runtime/device_pool.py index 794f6cd..d585a24 100644 --- a/artemis/runtime/device_pool.py +++ b/artemis/runtime/device_pool.py @@ -436,7 +436,15 @@ class DevicePool: selects the first ready device so the task will queue on it. Returns None if no devices are attached. """ - all_devs = self.list_devices() + return self._pick_device(self.list_devices(), preferred_serial) + + async def select_device_async(self, preferred_serial: str | None = None) -> str | None: + """Select a device without blocking the caller on ADB enumeration.""" + return self._pick_device(await self.list_devices_async(), preferred_serial) + + def _pick_device( + self, all_devs: list[DeviceStatus], preferred_serial: str | None + ) -> str | None: ready_devs = [d for d in all_devs if d.state == "device"] if preferred_serial: diff --git a/artemis/runtime/trace_store.py b/artemis/runtime/trace_store.py index 256fbab..09f99c8 100644 --- a/artemis/runtime/trace_store.py +++ b/artemis/runtime/trace_store.py @@ -326,6 +326,23 @@ def update_trace_status( return data +def update_trace_pid(trace_id: str, pid: int) -> dict[str, Any] | None: + """Updates the pid field of the status.json for a given trace_id.""" + path = get_status_path(trace_id) + with _status_lock(path): + data = read_status(trace_id) + if not data: + if os.path.exists(path) or os.path.exists(f"{path}.corrupt"): + logger.warning( + f"Dropping pid update for trace {trace_id}: existing status.json is unreadable" + ) + return None + + data["pid"] = pid + write_status(trace_id, data) + return data + + def update_trace_device_serial(trace_id: str, device_serial: str) -> dict[str, Any] | None: """Updates the device_serial field of the status.json for a given trace_id.""" path = get_status_path(trace_id) diff --git a/artemis/services/llm.py b/artemis/services/llm.py index 2c770ea..2a73cf9 100644 --- a/artemis/services/llm.py +++ b/artemis/services/llm.py @@ -21,7 +21,6 @@ thought stream recording, and role-based dynamic dispatching. import asyncio from collections.abc import Awaitable, Callable, Coroutine from contextvars import ContextVar, Token -from dataclasses import replace import functools import logging from pathlib import Path @@ -110,6 +109,9 @@ _ENDPOINT_BREAKER = CircuitBreaker(threshold=3, cooldown_seconds=30.0) # remembered so subsequent calls go straight to non-streaming invocation. _NON_STREAMING_ENDPOINTS: set[str] = set() +# Total retry limit when failures switch categories within one recovery cycle. +_MAX_TOTAL_RETRY_ATTEMPTS = 8 + def _get_current_data_engine(): if _CURRENT_DATA_ENGINE is not None: @@ -442,6 +444,13 @@ async def _run_with_recovery[T]( policy = retry_policy_for(failure.category) if attempt >= policy.max_attempts: break + total_attempts = sum(attempts.values()) + if total_attempts >= _MAX_TOTAL_RETRY_ATTEMPTS: + llm_logger.warning( + "LLM retry total across failure categories reached" + f" {total_attempts}; treating as exhausted." + ) + break delay = policy.delay_for(attempt) _record_llm_retry( str(e), @@ -753,10 +762,17 @@ class RobustChatModelWrapper: def _is_stream_unsupported_error(error: BaseException) -> bool: - """Detect 'this endpoint/model cannot stream' as opposed to a transient failure.""" - if isinstance(error, (NotImplementedError, AttributeError, TypeError)): + """Detect 'this endpoint/model cannot stream' as opposed to a transient failure. + + A positive here permanently disables streaming for the endpoint in this + process, so a bare AttributeError/TypeError from unrelated code must not + qualify — the message has to actually implicate streaming. + """ + if isinstance(error, NotImplementedError): return True message = str(error).lower() + if isinstance(error, (AttributeError, TypeError)): + return "stream" in message return "stream" in message and any( marker in message for marker in ("not support", "unsupported", "not implemented", "not available") @@ -1106,7 +1122,7 @@ def get_llm( """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) + endpoint = endpoint.model_copy(update={"temperature": temperature}) raw_model = ModelFactory.get_model(endpoint) handler = DataEngineCallbackHandler(ctx) diff --git a/artemis/tools/base.py b/artemis/tools/base.py index 6780b0d..74f2a25 100644 --- a/artemis/tools/base.py +++ b/artemis/tools/base.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Universal Tool Protocol and Registry for ARTEMIS. +"""Universal Tool Protocol 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 @@ -197,91 +197,3 @@ class ArtemisTool: 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, available_only: bool = True - ) -> list[ArtemisTool]: - """Lists all registered tools, optionally filtered by category and availability.""" - tools = [t for t in cls._tools.values() if not available_only or t.is_available()] - if category: - return [t for t in tools if t.category == category] - return list(tools) - - @classmethod - def get_langchain_tools( - cls, ctx: Any, names: list[str] | None = None, available_only: bool = True - ) -> 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()) - ) - if available_only: - tools = [t for t in tools if t.is_available(ctx)] - return [t.to_langchain_tool(ctx) for t in tools] - - @classmethod - def get_genai_declarations( - cls, names: list[str] | None = None, available_only: bool = True - ) -> 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()) - ) - if available_only: - tools = [t for t in tools if t.is_available()] - 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 index 0f522b6..360a36d 100644 --- a/artemis/tools/command_tool.py +++ b/artemis/tools/command_tool.py @@ -54,7 +54,7 @@ from artemis.controllers.platform_specific_commands_controller import ( ) 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.base import ArtemisTool from artemis.tools.tool_wrapper import ToolWrapper from artemis.utils.logger import get_logger @@ -832,7 +832,6 @@ class RunAdbCommandTool(ArtemisTool): # Universal tool instance & aliases run_adb_command = RunAdbCommandTool() RunAdbCommand = RunAdbCommandTool -ToolRegistry.register(run_adb_command) def get_run_adb_command_tool(ctx: ArtemisContext) -> BaseTool: @@ -978,7 +977,6 @@ class ManageTaskTool(ArtemisTool): # Universal tool instance & aliases manage_task = ManageTaskTool() ManageTask = ManageTaskTool -ToolRegistry.register(manage_task) def get_manage_task_tool(ctx: ArtemisContext) -> BaseTool: @@ -1111,7 +1109,6 @@ class RunShortAdbCommandTool(ArtemisTool): # 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: @@ -1198,7 +1195,6 @@ class AnalyzeTaskOutputTool(ArtemisTool): # Universal tool instance & aliases analyze_task_output = AnalyzeTaskOutputTool() AnalyzeTaskOutput = AnalyzeTaskOutputTool -ToolRegistry.register(analyze_task_output) def get_analyze_task_output_tool(ctx: ArtemisContext) -> BaseTool: diff --git a/artemis/tools/committee_tool.py b/artemis/tools/committee_tool.py index a964f72..9598814 100644 --- a/artemis/tools/committee_tool.py +++ b/artemis/tools/committee_tool.py @@ -36,7 +36,7 @@ from artemis.data_engine.trace import ( 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.base import ArtemisTool, ToolCategory 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 ( @@ -108,7 +108,6 @@ class AskCommitteeTool(ArtemisTool): ask_committee = AskCommitteeTool() AskCommittee = AskCommitteeTool AskCommitteeToolAlias = AskCommitteeTool -ToolRegistry.register(ask_committee) def get_ask_committee_tool(ctx: ArtemisContext) -> BaseTool: diff --git a/artemis/tools/diagnoser_submit_answer_tool.py b/artemis/tools/diagnoser_submit_answer_tool.py index ee79f69..abdb7f3 100644 --- a/artemis/tools/diagnoser_submit_answer_tool.py +++ b/artemis/tools/diagnoser_submit_answer_tool.py @@ -22,7 +22,7 @@ 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.tools.base import ArtemisTool from artemis.utils.logger import get_logger logger = get_logger(__name__) @@ -80,7 +80,6 @@ class SubmitAnswerTool(ArtemisTool): submit_answer = SubmitAnswerTool() SubmitAnswer = SubmitAnswerTool DiagnoserSubmitAnswerTool = SubmitAnswerTool -ToolRegistry.register(submit_answer) def get_submit_answer_tool(ctx: ArtemisContext) -> BaseTool: diff --git a/artemis/tools/diagnostic_tool.py b/artemis/tools/diagnostic_tool.py index 8c7eb13..90dbe0b 100644 --- a/artemis/tools/diagnostic_tool.py +++ b/artemis/tools/diagnostic_tool.py @@ -24,7 +24,7 @@ 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.base import ArtemisTool from artemis.tools.tool_wrapper import ToolWrapper from artemis.tools.types import CyFunctionDetector from artemis.utils.logger import get_logger @@ -100,7 +100,6 @@ class AskDiagnoserTool(ArtemisTool): ask_diagnoser = AskDiagnoserTool() AskDiagnoser = AskDiagnoserTool DiagnosticTool = AskDiagnoserTool -ToolRegistry.register(ask_diagnoser) def get_ask_diagnostic_agent_tool(ctx: ArtemisContext) -> BaseTool: diff --git a/artemis/tools/explorer_tool.py b/artemis/tools/explorer_tool.py index 311a2fc..9065e0e 100644 --- a/artemis/tools/explorer_tool.py +++ b/artemis/tools/explorer_tool.py @@ -56,7 +56,7 @@ from artemis.config import ExplorerVersion, resolve_explorer_version, settings from artemis.context import ArtemisContext from artemis.drivers.base import BaseDeviceDriver from artemis.graph.state import State -from artemis.tools.base import ArtemisTool, ToolCategory, ToolRegistry +from artemis.tools.base import ArtemisTool, ToolCategory from artemis.tools.tool_wrapper import ToolWrapper from artemis.tools.types import CyFunctionDetector from artemis.utils.element_hit_test import find_element_at_point @@ -527,7 +527,6 @@ class AskExplorerTool(ArtemisTool): ask_explorer = AskExplorerTool() -ToolRegistry.register(ask_explorer) def get_ask_explorer_tool( diff --git a/artemis/tools/history_recall.py b/artemis/tools/history_recall.py index 9afea9f..389219e 100644 --- a/artemis/tools/history_recall.py +++ b/artemis/tools/history_recall.py @@ -45,7 +45,7 @@ from pydantic import BaseModel, Field from artemis.core.tool_declaration import ToolDeclaration from artemis.data_engine.trace import trace_langchain_tool -from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.base import ArtemisTool from artemis.tools.tool_wrapper import ToolWrapper from artemis.utils.logger import get_logger @@ -547,7 +547,6 @@ def _recall_config(warn: bool = True) -> Any: recall_history = RecallHistoryTool() -ToolRegistry.register(recall_history) _DECLARATION_TYPES = { "query": {"type": "string"}, diff --git a/artemis/tools/image_processor_tool.py b/artemis/tools/image_processor_tool.py index 2df9ef4..33b22c3 100644 --- a/artemis/tools/image_processor_tool.py +++ b/artemis/tools/image_processor_tool.py @@ -23,7 +23,7 @@ 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.base import ArtemisTool from artemis.tools.tool_wrapper import ToolWrapper from artemis.tools.types import CyFunctionDetector from artemis.utils.logger import get_logger @@ -90,7 +90,6 @@ ask_image_processor = AskImageProcessorTool() AskImageProcessor = AskImageProcessorTool AskVisionCoderTool = AskImageProcessorTool ImageProcessorTool = AskImageProcessorTool -ToolRegistry.register(ask_image_processor) def get_ask_image_processor_tool(ctx: ArtemisContext) -> BaseTool: diff --git a/artemis/tools/log_tool.py b/artemis/tools/log_tool.py index 13add08..7b3525c 100644 --- a/artemis/tools/log_tool.py +++ b/artemis/tools/log_tool.py @@ -24,7 +24,7 @@ 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.base import ArtemisTool from artemis.tools.tool_wrapper import ToolWrapper from artemis.tools.types import CyFunctionDetector from artemis.utils.logger import get_logger @@ -93,7 +93,6 @@ class AnalyzeLogsTool(ArtemisTool): analyze_logs = AnalyzeLogsTool() AnalyzeLogs = AnalyzeLogsTool LogTool = AnalyzeLogsTool -ToolRegistry.register(analyze_logs) def get_analyze_logs_tool(ctx: ArtemisContext) -> BaseTool: diff --git a/artemis/tools/memory/notes.py b/artemis/tools/memory/notes.py deleted file mode 100644 index 5f81ef0..0000000 --- a/artemis/tools/memory/notes.py +++ /dev/null @@ -1,66 +0,0 @@ -# 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/launch_app.py b/artemis/tools/mobile/launch_app.py index 999de34..8f4cdf1 100644 --- a/artemis/tools/mobile/launch_app.py +++ b/artemis/tools/mobile/launch_app.py @@ -26,7 +26,7 @@ from artemis.controllers.platform_specific_commands_controller import ( 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.base import ArtemisTool, ToolCategory from artemis.tools.tool_wrapper import ToolWrapper from artemis.tools.types import CyFunctionDetector from artemis.utils.app_launch_utils import launch_app_with_retries @@ -189,7 +189,6 @@ class LaunchAppTool(ArtemisTool): # Universal tool instance & aliases launch_app = LaunchAppTool() LaunchApp = LaunchAppTool -ToolRegistry.register(launch_app) def get_launch_app_tool(ctx: ArtemisContext) -> BaseTool: diff --git a/artemis/tools/mobile/ocr.py b/artemis/tools/mobile/ocr.py index 866379c..a3f3ffc 100644 --- a/artemis/tools/mobile/ocr.py +++ b/artemis/tools/mobile/ocr.py @@ -25,7 +25,7 @@ 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.base import ArtemisTool, ToolCategory from artemis.tools.tool_wrapper import ToolWrapper from artemis.utils.logger import get_logger from artemis.utils.ocr_api import is_ocr_configured, perform_ocr @@ -187,7 +187,6 @@ ocr_recognition = OcrRecognitionTool() OcrRecognition = OcrRecognitionTool OCRRecognition = OcrRecognitionTool OcrTool = OcrRecognitionTool -ToolRegistry.register(ocr_recognition) def get_ocr_tool(ctx: ArtemisContext) -> BaseTool: diff --git a/artemis/tools/mobile/read_hierarchy.py b/artemis/tools/mobile/read_hierarchy.py index f05d4ed..35f9144 100644 --- a/artemis/tools/mobile/read_hierarchy.py +++ b/artemis/tools/mobile/read_hierarchy.py @@ -22,7 +22,7 @@ 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.base import ArtemisTool, ToolCategory from artemis.tools.tool_wrapper import ToolWrapper from artemis.tools.types import CyFunctionDetector from artemis.utils.logger import get_logger @@ -91,7 +91,6 @@ class GetUiHierarchyTool(ArtemisTool): get_ui_hierarchy = GetUiHierarchyTool() GetUiHierarchy = GetUiHierarchyTool GetUIHierarchy = GetUiHierarchyTool -ToolRegistry.register(get_ui_hierarchy) def get_ui_hierarchy_tool(ctx: ArtemisContext) -> BaseTool: diff --git a/artemis/tools/mobile/read_logs.py b/artemis/tools/mobile/read_logs.py index bfab9de..5e5bb57 100644 --- a/artemis/tools/mobile/read_logs.py +++ b/artemis/tools/mobile/read_logs.py @@ -21,7 +21,7 @@ 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.base import ArtemisTool, ToolCategory from artemis.tools.mobile.log_utils import fetch_and_filter_logs from artemis.utils.logger import get_logger @@ -108,7 +108,6 @@ class ReadLogsTool(ArtemisTool): read_logs = ReadLogsTool() ReadLogs = ReadLogsTool ReadLogsToolAlias = ReadLogsTool -ToolRegistry.register(read_logs) def get_read_logs_tool(ctx: ArtemisContext) -> BaseTool: diff --git a/artemis/tools/mobile/search_logs.py b/artemis/tools/mobile/search_logs.py index af328ef..42307a2 100644 --- a/artemis/tools/mobile/search_logs.py +++ b/artemis/tools/mobile/search_logs.py @@ -22,7 +22,7 @@ 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.base import ArtemisTool, ToolCategory from artemis.tools.mobile.log_utils import fetch_and_filter_logs from artemis.utils.logger import get_logger @@ -229,7 +229,6 @@ class SearchLogsTool(ArtemisTool): search_logs = SearchLogsTool() SearchLogs = SearchLogsTool SearchLogsToolAlias = SearchLogsTool -ToolRegistry.register(search_logs) def get_search_logs_tool(ctx: ArtemisContext) -> BaseTool: diff --git a/artemis/tools/object_detection_tool.py b/artemis/tools/object_detection_tool.py index e0351c6..c922f90 100644 --- a/artemis/tools/object_detection_tool.py +++ b/artemis/tools/object_detection_tool.py @@ -33,7 +33,7 @@ from artemis.data_engine.trace import ( ) from artemis.drivers.base import BaseDeviceDriver from artemis.graph.state import State -from artemis.tools.base import ArtemisTool, ToolCategory, ToolRegistry +from artemis.tools.base import ArtemisTool, ToolCategory from artemis.tools.tool_wrapper import ToolWrapper from artemis.tools.types import CyFunctionDetector from artemis.utils.logger import get_logger @@ -220,8 +220,6 @@ 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.""" diff --git a/artemis/tools/perception/object_detect.py b/artemis/tools/perception/object_detect.py deleted file mode 100644 index 76716ea..0000000 --- a/artemis/tools/perception/object_detect.py +++ /dev/null @@ -1,43 +0,0 @@ -# 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 deleted file mode 100644 index e88c7b7..0000000 --- a/artemis/tools/perception/ocr.py +++ /dev/null @@ -1,46 +0,0 @@ -# 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 deleted file mode 100644 index cb13c97..0000000 --- a/artemis/tools/perception/xml_search.py +++ /dev/null @@ -1,45 +0,0 @@ -# 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 index 35b65b0..34f2fba 100644 --- a/artemis/tools/scratchpad.py +++ b/artemis/tools/scratchpad.py @@ -22,7 +22,7 @@ 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.base import ArtemisTool from artemis.tools.tool_wrapper import ToolWrapper from artemis.tools.types import CyFunctionDetector from artemis.utils.logger import get_logger @@ -166,7 +166,6 @@ class SaveNoteTool(ArtemisTool): # Universal tool instance & aliases save_note = SaveNoteTool() SaveNote = SaveNoteTool -ToolRegistry.register(save_note) def get_save_note_tool(ctx: ArtemisContext) -> BaseTool: @@ -246,7 +245,6 @@ class AppendNoteTool(ArtemisTool): # Universal tool instance & aliases append_note = AppendNoteTool() AppendNote = AppendNoteTool -ToolRegistry.register(append_note) def get_append_note_tool(ctx: ArtemisContext) -> BaseTool: @@ -331,7 +329,6 @@ class UpdateNoteTool(ArtemisTool): # Universal tool instance & aliases update_note = UpdateNoteTool() UpdateNote = UpdateNoteTool -ToolRegistry.register(update_note) def get_update_note_tool(ctx: ArtemisContext) -> BaseTool: @@ -416,7 +413,6 @@ class ReadNoteTool(ArtemisTool): # Universal tool instance & aliases read_note = ReadNoteTool() ReadNote = ReadNoteTool -ToolRegistry.register(read_note) def get_read_note_tool(ctx: ArtemisContext) -> BaseTool: @@ -478,7 +474,6 @@ class ListNotesTool(ArtemisTool): # Universal tool instance & aliases list_notes = ListNotesTool() ListNotes = ListNotesTool -ToolRegistry.register(list_notes) def get_list_notes_tool(ctx: ArtemisContext) -> BaseTool: @@ -545,7 +540,6 @@ class ReadNotePureTool(ArtemisTool): # 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: @@ -591,7 +585,6 @@ class ListNotesPureTool(ArtemisTool): # 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: @@ -656,7 +649,6 @@ class SaveNotePureTool(ArtemisTool): # 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: @@ -725,7 +717,6 @@ class UpdateNotePureTool(ArtemisTool): # 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: @@ -790,7 +781,6 @@ class AppendNotePureTool(ArtemisTool): # 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: diff --git a/artemis/tools/tool_wrapper.py b/artemis/tools/tool_wrapper.py index 073acd1..b455504 100644 --- a/artemis/tools/tool_wrapper.py +++ b/artemis/tools/tool_wrapper.py @@ -13,10 +13,10 @@ # limitations under the License. from collections.abc import Callable +import functools import inspect import time from typing import Annotated, Any, get_args, get_origin -from unittest.mock import MagicMock, Mock import uuid from langchain_core.messages import ToolMessage @@ -64,27 +64,19 @@ async def invoke_tool_with_injection( """ 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) + def _real_callable(candidate: Any) -> Callable | None: + """Return a routine or partial whose signature supports direct injection.""" + target = candidate.func if isinstance(candidate, functools.partial) else candidate + if inspect.isroutine(target): + return candidate + return None # 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 getattr(tool, "coroutine", None): + func = _real_callable(tool.coroutine) + if func is None and getattr(tool, "func", None): + func = _real_callable(tool.func) if func: try: @@ -163,7 +155,11 @@ async def invoke_tool_with_injection( ) CURRENT_TRACE_ID.reset(token) else: - # Fallback to ainvoke if no direct func available + # Fallback to ainvoke if no direct func available; tools that declare + # a "state" argument still receive the graph state. + tool_args = getattr(tool, "args", None) + if state is not None and isinstance(tool_args, dict) and "state" in tool_args: + final_args["state"] = state return await tool.ainvoke(final_args) diff --git a/artemis/tools/video_tool.py b/artemis/tools/video_tool.py index 93984e7..7842837 100644 --- a/artemis/tools/video_tool.py +++ b/artemis/tools/video_tool.py @@ -31,7 +31,7 @@ from artemis.core.tool_declaration import ToolDeclaration 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.base import ArtemisTool, ToolCategory from artemis.tools.tool_wrapper import ToolWrapper from artemis.tools.types import CyFunctionDetector from artemis.utils.logger import get_logger @@ -180,7 +180,6 @@ class VideoAnalyzerTool(ArtemisTool): # Universal tool instance & aliases video_analyzer = VideoAnalyzerTool() -ToolRegistry.register(video_analyzer) class VideoAnalyzerPureTool(ArtemisTool): @@ -231,7 +230,6 @@ class VideoAnalyzerPureTool(ArtemisTool): # Universal pure tool instance & aliases video_analyzer_pure = VideoAnalyzerPureTool() VideoAnalyzerPure = VideoAnalyzerPureTool -ToolRegistry.register(video_analyzer_pure) def get_video_analyzer_tool( diff --git a/artemis/tools/wait_tool.py b/artemis/tools/wait_tool.py index d3662da..866a573 100644 --- a/artemis/tools/wait_tool.py +++ b/artemis/tools/wait_tool.py @@ -21,7 +21,7 @@ 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.base import ArtemisTool, ToolCategory from artemis.tools.tool_wrapper import ToolWrapper from artemis.tools.types import CyFunctionDetector from artemis.utils.logger import get_logger @@ -82,7 +82,6 @@ class WaitTool(ArtemisTool): # Universal tool instance & aliases wait = WaitTool() Wait = WaitTool -ToolRegistry.register(wait) def get_wait_tool(ctx: ArtemisContext) -> BaseTool: diff --git a/tests/fixtures/action_surfaces/tool_declarations.json b/tests/fixtures/action_surfaces/tool_declarations.json index 394ef97..a8dd936 100644 --- a/tests/fixtures/action_surfaces/tool_declarations.json +++ b/tests/fixtures/action_surfaces/tool_declarations.json @@ -27,31 +27,6 @@ "target" ] } - }, - "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": { @@ -82,31 +57,6 @@ "sequence" ] } - }, - "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": { @@ -133,27 +83,6 @@ "target" ] } - }, - "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": { @@ -185,32 +114,6 @@ "target" ] } - }, - "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": { @@ -257,47 +160,6 @@ } } } - }, - "name": "swipe", - "description": "[ACTION] Perform a swipe, drag, or slider-adjustment gesture on the screen. The screen and UI hierarchy after swipe will be returned automatically.\n\n• Directional Scrolling ('direction' or 'action'): Recommended for general browsing and standard page scrolling in most scenarios. Automatically computes safe swipe vectors and adaptive duration, retains a ~40% visual overlap anchor for zero-omission traversal, and prevents inertial flings. Supports scoping to a sub-container via 'target'. If it fails on certain custom layouts, fall back to specifying exact coordinates ('start' and 'end') directly.\n• Precise Coordinate Gestures ('start', 'end' or coordinates list): Best for local, fine-grained interactions such as adjusting sliders/SeekBars (e.g., volume, brightness, progress bars), drag-and-drop / list reordering, or as a reliable fallback when directional scrolling fails on specific containers. 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.", - "parameters": { - "type": "object", - "properties": { - "direction": { - "type": "string", - "enum": [ - "up", - "down", - "left", - "right" - ], - "description": "Direction for scrolling and swiping: '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": { - "type": "array", - "items": { - "type": "integer" - }, - "description": "Start normalized coordinates [start_x, start_y] in 0-1000 scale." - }, - "end": { - "type": "array", - "items": { - "type": "integer" - }, - "description": "End normalized coordinates [end_x, end_y] in 0-1000 scale." - }, - "target": { - "description": "Optional target element index (e.g. 2) or container bounds [left, top, right, bottom] to scope the directional swipe within." - }, - "action": { - "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": { - "type": "integer", - "description": "Optional swipe/drag duration in milliseconds (default 800). 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." - } - } } }, "PRESS_KEY_TOOL": { @@ -317,20 +179,6 @@ "key" ] } - }, - "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" - ] } }, "MANAGE_APP_TOOL": { @@ -355,25 +203,6 @@ "app_name" ] } - }, - "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": { @@ -393,20 +222,6 @@ "time_in_ms" ] } - }, - "name": "wait_for_delay", - "description": "[ACTION] Pause execution and wait for a specified 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 for 2s, 60000 for 1 min, 180000 for 3 mins, 300000 for 5 mins). Convert any required waiting duration into milliseconds." - } - }, - "required": [ - "time_in_ms" - ] } } } diff --git a/tests/integration/agents/explorer/test_explorer_all_tools/test_explorer_all_tools.py b/tests/integration/agents/explorer/test_explorer_all_tools/test_explorer_all_tools.py index 111108c..17682c2 100644 --- a/tests/integration/agents/explorer/test_explorer_all_tools/test_explorer_all_tools.py +++ b/tests/integration/agents/explorer/test_explorer_all_tools/test_explorer_all_tools.py @@ -87,7 +87,7 @@ async def test_explorer_all_tools_sequential_mocked(): logger.addHandler(handler) # Attach the main agent logger to capture internal step logs - agent_logger = logging.getLogger("artemis.agents.explorer.explorer") + agent_logger = logging.getLogger("artemis.agents.explorer") agent_logger.setLevel(logging.INFO) agent_logger.addHandler(handler) @@ -238,12 +238,13 @@ async def test_explorer_all_tools_sequential_mocked(): return_value=mock_client, ), patch( - "artemis.agents.explorer.explorer.StorageManager", + "artemis.agents.explorer.run_setup.StorageManager", return_value=mock_storage, ), # get_ocr_list is only exposed when an OCR backend is configured. patch("artemis.agents.explorer.explorer.is_ocr_configured", return_value=True), - patch("artemis.agents.explorer.explorer._run_object_detection", mock_detect), + patch("artemis.agents.explorer.run_setup.is_ocr_configured", return_value=True), + patch("artemis.agents.explorer.perception_tools._run_object_detection", mock_detect), ): explorer = Explorer(mock_ctx) result = await explorer.run( diff --git a/tests/integration/agents/explorer/test_explorer_all_tools/test_explorer_ask_image_processor.py b/tests/integration/agents/explorer/test_explorer_all_tools/test_explorer_ask_image_processor.py index efe4105..ce7ef00 100644 --- a/tests/integration/agents/explorer/test_explorer_all_tools/test_explorer_ask_image_processor.py +++ b/tests/integration/agents/explorer/test_explorer_all_tools/test_explorer_ask_image_processor.py @@ -84,7 +84,7 @@ async def test_explorer_ask_image_processor_tool(): logger.addHandler(handler) # Attach the main agent logger to capture internal step logs - agent_logger = logging.getLogger("artemis.agents.explorer.explorer") + agent_logger = logging.getLogger("artemis.agents.explorer") agent_logger.setLevel(logging.INFO) agent_logger.addHandler(handler) diff --git a/tests/tools/diagnoser/test_diagnoser_tools.py b/tests/tools/diagnoser/test_diagnoser_tools.py index 554645a..fbc9574 100644 --- a/tests/tools/diagnoser/test_diagnoser_tools.py +++ b/tests/tools/diagnoser/test_diagnoser_tools.py @@ -31,9 +31,9 @@ def test_diagnoser_submit_answer_tool(artemis_context): assert "error" not in result.lower() -def test_submit_answer_tool_subclass_and_registry(): - """Verify SubmitAnswerTool is an ArtemisTool subclass and registered in ToolRegistry.""" - from artemis.tools.base import ArtemisTool, ToolRegistry +def test_submit_answer_tool_subclass(): + """Verify SubmitAnswerTool is an ArtemisTool subclass.""" + from artemis.tools.base import ArtemisTool from artemis.tools.diagnoser_submit_answer_tool import ( DiagnoserSubmitAnswerTool, SubmitAnswer, @@ -52,11 +52,6 @@ def test_submit_answer_tool_subclass_and_registry(): assert submit_answer.category == "custom" assert submit_answer.args_schema == SubmitAnswerArgs - # Registry lookup - reg_tool = ToolRegistry.get("submit_answer") - assert reg_tool is not None - assert isinstance(reg_tool, SubmitAnswerTool) - # GenAI FunctionDeclaration export declaration = submit_answer.to_genai_declaration() assert declaration.name == "submit_answer" @@ -64,9 +59,9 @@ def test_submit_answer_tool_subclass_and_registry(): assert "actionable_steps" in declaration.parameters.properties -def test_ask_diagnoser_tool_subclass_and_registry(): - """Verify AskDiagnoserTool is an ArtemisTool subclass and registered in ToolRegistry.""" - from artemis.tools.base import ArtemisTool, ToolRegistry +def test_ask_diagnoser_tool_subclass(): + """Verify AskDiagnoserTool is an ArtemisTool subclass.""" + from artemis.tools.base import ArtemisTool from artemis.tools.diagnostic_tool import ( AskDiagnoser, AskDiagnoserArgs, @@ -85,20 +80,15 @@ def test_ask_diagnoser_tool_subclass_and_registry(): assert ask_diagnoser.category == "custom" assert ask_diagnoser.args_schema == AskDiagnoserArgs - # Registry lookup - reg_tool = ToolRegistry.get("ask_diagnoser") - assert reg_tool is not None - assert isinstance(reg_tool, AskDiagnoserTool) - # GenAI FunctionDeclaration export declaration = ask_diagnoser.to_genai_declaration() assert declaration.name == "ask_diagnoser" assert "query" in declaration.parameters.properties -def test_analyze_logs_tool_subclass_and_registry(): - """Verify AnalyzeLogsTool is an ArtemisTool subclass and registered in ToolRegistry.""" - from artemis.tools.base import ArtemisTool, ToolRegistry +def test_analyze_logs_tool_subclass(): + """Verify AnalyzeLogsTool is an ArtemisTool subclass.""" + from artemis.tools.base import ArtemisTool from artemis.tools.log_tool import ( AnalyzeLogs, AnalyzeLogsArgs, @@ -117,11 +107,6 @@ def test_analyze_logs_tool_subclass_and_registry(): assert analyze_logs.category == "custom" assert analyze_logs.args_schema == AnalyzeLogsArgs - # Registry lookup - reg_tool = ToolRegistry.get("analyze_logs") - assert reg_tool is not None - assert isinstance(reg_tool, AnalyzeLogsTool) - # GenAI FunctionDeclaration export declaration = analyze_logs.to_genai_declaration() assert declaration.name == "analyze_logs" diff --git a/tests/tools/object_detector/test_object_detector_tools.py b/tests/tools/object_detector/test_object_detector_tools.py index 48ffd63..b3d3493 100644 --- a/tests/tools/object_detector/test_object_detector_tools.py +++ b/tests/tools/object_detector/test_object_detector_tools.py @@ -14,7 +14,7 @@ """Tests for object_detector universal tools.""" -from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.base import ArtemisTool from artemis.tools.object_detection_tool import ( ObjectDetection, ObjectDetectionArgs, @@ -29,8 +29,8 @@ from artemis.tools.object_detection_tool import ( ) -def test_object_detector_tool_subclass_and_registry(): - """Verify ObjectDetectionTool is an ArtemisTool subclass and registered in ToolRegistry.""" +def test_object_detector_tool_subclass(): + """Verify ObjectDetectionTool is an ArtemisTool subclass.""" assert issubclass(ObjectDetectionTool, ArtemisTool) assert issubclass(ObjectDetection, ArtemisTool) assert issubclass(ObjectDetectorTool, ArtemisTool) @@ -51,8 +51,3 @@ def test_object_detector_tool_subclass_and_registry(): assert operator_object_detection.name == "object_detection" assert operator_object_detection.category == "perception" assert operator_object_detection.args_schema == OperatorObjectDetectionArgs - - # Registry lookup - reg_tool = ToolRegistry.get("object_detection") - assert reg_tool is not None - assert isinstance(reg_tool, ObjectDetectionTool) diff --git a/tests/unit/admin_console/test_task_queue_service.py b/tests/unit/admin_console/test_task_queue_service.py index 158723b..24a572c 100644 --- a/tests/unit/admin_console/test_task_queue_service.py +++ b/tests/unit/admin_console/test_task_queue_service.py @@ -112,7 +112,7 @@ async def test_enqueued_task_keeps_its_adb_endpoint_snapshot(): return_value=original, ), patch( - "artemis.runtime.device_pool.device_pool.select_device", + "artemis.runtime.device_pool.device_pool.select_device_async", return_value="emulator-5554", ), ): @@ -230,7 +230,7 @@ async def test_queue_worker_execution_lifecycle(): patch("apps.admin_console.services.task_queue_service.session_repo") as mock_repo, patch("apps.admin_console.services.task_queue_service.media_service"), patch( - "artemis.runtime.device_pool.device_pool.select_device", + "artemis.runtime.device_pool.device_pool.select_device_async", return_value="emulator-5554", ), ): @@ -276,7 +276,7 @@ async def test_queue_worker_cmd_construction(): patch("apps.admin_console.services.task_queue_service.session_repo") as mock_repo, patch("apps.admin_console.services.task_queue_service.media_service"), patch( - "artemis.runtime.device_pool.device_pool.select_device", + "artemis.runtime.device_pool.device_pool.select_device_async", return_value="emulator-5554", ), ): @@ -594,7 +594,7 @@ async def test_cancel_task_triggers_next_pending_task(): patch("apps.admin_console.services.task_queue_service.session_repo") as mock_repo, patch("apps.admin_console.services.task_queue_service.media_service"), patch( - "artemis.runtime.device_pool.device_pool.select_device", + "artemis.runtime.device_pool.device_pool.select_device_async", return_value="emulator-5554", ), ): @@ -645,7 +645,7 @@ async def test_immediate_cancel_ignores_stale_ipc_and_runs_next_task(): patch("apps.admin_console.services.task_queue_service.session_repo") as mock_repo, patch("apps.admin_console.services.task_queue_service.media_service"), patch( - "artemis.runtime.device_pool.device_pool.select_device", + "artemis.runtime.device_pool.device_pool.select_device_async", return_value="emulator-5554", ), ): diff --git a/tests/unit/agents/test_explorer.py b/tests/unit/agents/test_explorer.py index 21af249..0735746 100644 --- a/tests/unit/agents/test_explorer.py +++ b/tests/unit/agents/test_explorer.py @@ -106,7 +106,7 @@ async def test_explorer_run(): return_value=mock_client, ), patch( - "artemis.agents.explorer.explorer.StorageManager", + "artemis.agents.explorer.run_setup.StorageManager", return_value=mock_storage, ), patch("pathlib.Path.exists", return_value=True), @@ -193,7 +193,7 @@ async def test_explorer_submit_answer(): return_value=mock_client, ), patch( - "artemis.agents.explorer.explorer.StorageManager", + "artemis.agents.explorer.run_setup.StorageManager", return_value=mock_storage, ), patch("pathlib.Path.exists", return_value=True), @@ -339,7 +339,7 @@ async def test_explorer_submit_answer_self_correction(): return_value=mock_client, ), patch( - "artemis.agents.explorer.explorer.StorageManager", + "artemis.agents.explorer.run_setup.StorageManager", return_value=mock_storage, ), patch("pathlib.Path.exists", return_value=True), @@ -427,10 +427,10 @@ async def test_explorer_initial_visual_marking(): return_value=mock_client, ), patch( - "artemis.agents.explorer.explorer.StorageManager", + "artemis.agents.explorer.run_setup.StorageManager", return_value=mock_storage, ), - patch("artemis.agents.explorer.explorer.draw_dots") as mock_draw_dots, + patch("artemis.agents.explorer.run_setup.draw_dots") as mock_draw_dots, patch("pathlib.Path.exists", return_value=True), patch("pathlib.Path.read_text", return_value=mock_prompt), patch("builtins.open", custom_open), @@ -518,10 +518,10 @@ async def test_explorer_initial_visual_marking_previous_screenshot(): return_value=mock_client, ), patch( - "artemis.agents.explorer.explorer.StorageManager", + "artemis.agents.explorer.run_setup.StorageManager", return_value=mock_storage, ), - patch("artemis.agents.explorer.explorer.draw_dots") as mock_draw_dots, + patch("artemis.agents.explorer.run_setup.draw_dots") as mock_draw_dots, patch("pathlib.Path.exists", return_value=True), patch("pathlib.Path.read_text", return_value=mock_prompt), patch("builtins.open", custom_open), @@ -606,10 +606,10 @@ async def test_explorer_initial_visual_marking_previous_screenshot_no_ui_tree(): return_value=mock_client, ), patch( - "artemis.agents.explorer.explorer.StorageManager", + "artemis.agents.explorer.run_setup.StorageManager", return_value=mock_storage, ), - patch("artemis.agents.explorer.explorer.draw_dots") as mock_draw_dots, + patch("artemis.agents.explorer.run_setup.draw_dots") as mock_draw_dots, patch("pathlib.Path.exists", return_value=True), patch("pathlib.Path.read_text", return_value=mock_prompt), patch("builtins.open", custom_open), @@ -704,10 +704,10 @@ async def test_explorer_initial_visual_marking_previous_screenshot_ocr_fusion(): return_value=mock_client, ), patch( - "artemis.agents.explorer.explorer.StorageManager", + "artemis.agents.explorer.run_setup.StorageManager", return_value=mock_storage, ), - patch("artemis.agents.explorer.explorer.draw_dots") as mock_draw_dots, + patch("artemis.agents.explorer.run_setup.draw_dots") as mock_draw_dots, patch("pathlib.Path.exists", return_value=True), patch("pathlib.Path.read_text", return_value=mock_prompt), patch("builtins.open", custom_open), @@ -808,19 +808,19 @@ async def test_explorer_initial_visual_marking_previous_screenshot_on_the_fly_oc return_value=mock_client, ), patch( - "artemis.agents.explorer.explorer.StorageManager", + "artemis.agents.explorer.run_setup.StorageManager", return_value=mock_storage, ), patch( - "artemis.agents.explorer.explorer.perform_ocr", + "artemis.agents.explorer.run_setup.perform_ocr", new_callable=AsyncMock, return_value=mock_ocr_result, ) as mock_perform_ocr, patch( - "artemis.agents.explorer.explorer.is_ocr_configured", + "artemis.agents.explorer.run_setup.is_ocr_configured", return_value=True, ), - patch("artemis.agents.explorer.explorer.draw_dots") as mock_draw_dots, + patch("artemis.agents.explorer.run_setup.draw_dots") as mock_draw_dots, patch("pathlib.Path.exists", return_value=True), patch("pathlib.Path.read_text", return_value=mock_prompt), patch("builtins.open", custom_open), @@ -933,7 +933,7 @@ async def test_explorer_denylisted_tool(): return_value=mock_client, ), patch( - "artemis.agents.explorer.explorer.StorageManager", + "artemis.agents.explorer.run_setup.StorageManager", return_value=mock_storage, ), patch("pathlib.Path.exists", return_value=True), @@ -1014,7 +1014,7 @@ async def test_explorer_inspect_region(): patch("cv2.imwrite", return_value=True) as mock_imwrite, patch("pathlib.Path.mkdir") as mock_mkdir, patch("glob.glob", return_value=[]), - patch("artemis.agents.explorer.explorer.settings") as mock_settings, + patch("artemis.agents.explorer.perception_tools.settings") as mock_settings, ): mock_settings.TRACES_PATH = pathlib.Path("/tmp/traces") res = await explorer.exec_inspect_region( @@ -1181,12 +1181,12 @@ async def test_explorer_ask_perception_tool(): mock_draw_dots = MagicMock() with ( - patch("artemis.agents.explorer.explorer.StorageManager") as storage_cls, + patch("artemis.agents.explorer.run_setup.StorageManager") as storage_cls, patch( - "artemis.agents.explorer.explorer._run_object_detection", + "artemis.agents.explorer.perception_tools._run_object_detection", return_value=mock_detector_res, ), - patch("artemis.agents.explorer.explorer.draw_dots", mock_draw_dots), + patch("artemis.agents.explorer.perception_tools.draw_dots", mock_draw_dots), patch("artemis.utils.visualization.draw_dots", mock_draw_dots), patch("pathlib.Path.mkdir"), patch("glob.glob", return_value=[]), @@ -1277,7 +1277,7 @@ class _FakeGenaiError(Exception): @pytest.mark.asyncio async def test_generate_content_reliability_permanent_error_is_not_retried(): - from artemis.agents.explorer.explorer import _generate_content_with_reliability + from artemis.agents.explorer.native_runner import _generate_content_with_reliability from artemis.llm.reliability import LLMPermanentError calls = 0 @@ -1296,7 +1296,7 @@ async def test_generate_content_reliability_permanent_error_is_not_retried(): @pytest.mark.asyncio async def test_generate_content_reliability_bad_request_is_not_retried(): - from artemis.agents.explorer.explorer import _generate_content_with_reliability + from artemis.agents.explorer.native_runner import _generate_content_with_reliability from artemis.llm.reliability import LLMPermanentError calls = 0 @@ -1314,7 +1314,7 @@ async def test_generate_content_reliability_bad_request_is_not_retried(): @pytest.mark.asyncio async def test_generate_content_reliability_retries_transient_then_succeeds(): - from artemis.agents.explorer import explorer as explorer_module + from artemis.agents.explorer import native_runner as native_runner_module calls = 0 sentinel = object() @@ -1326,8 +1326,8 @@ async def test_generate_content_reliability_retries_transient_then_succeeds(): raise _FakeGenaiError("429 rate limit exceeded", code=429) return sentinel - with patch.object(explorer_module.asyncio, "sleep", new=AsyncMock()) as sleep_mock: - result = await explorer_module._generate_content_with_reliability(op) + with patch.object(native_runner_module.asyncio, "sleep", new=AsyncMock()) as sleep_mock: + result = await native_runner_module._generate_content_with_reliability(op) assert result is sentinel assert calls == 2 assert sleep_mock.await_count == 1 @@ -1335,7 +1335,7 @@ async def test_generate_content_reliability_retries_transient_then_succeeds(): @pytest.mark.asyncio async def test_generate_content_reliability_exhausts_to_typed_error(): - from artemis.agents.explorer import explorer as explorer_module + from artemis.agents.explorer import native_runner as native_runner_module from artemis.llm.reliability import ( FailureCategory, LLMExhaustedError, @@ -1349,9 +1349,9 @@ async def test_generate_content_reliability_exhausts_to_typed_error(): calls += 1 raise _FakeGenaiError("503 service unavailable", code=503) - with patch.object(explorer_module.asyncio, "sleep", new=AsyncMock()): + with patch.object(native_runner_module.asyncio, "sleep", new=AsyncMock()): with pytest.raises(LLMExhaustedError) as exc_info: - await explorer_module._generate_content_with_reliability(op) + await native_runner_module._generate_content_with_reliability(op) assert calls == retry_policy_for(FailureCategory.PROVIDER_UNAVAILABLE).max_attempts assert exc_info.value.failure.category.value == "provider_unavailable" assert isinstance(exc_info.value.__cause__, _FakeGenaiError) diff --git a/tests/unit/agents/test_explorer_perception_regressions.py b/tests/unit/agents/test_explorer_perception_regressions.py index 18bdf18..559dfa1 100644 --- a/tests/unit/agents/test_explorer_perception_regressions.py +++ b/tests/unit/agents/test_explorer_perception_regressions.py @@ -37,7 +37,7 @@ def _perception_host() -> PerceptionToolsMixin: async def test_get_ocr_list_reads_the_screen_index_without_storage() -> None: host = _perception_host() - with patch("artemis.agents.explorer.explorer.StorageManager") as storage_cls: + with patch("artemis.agents.explorer.run_setup.StorageManager") as storage_cls: result = await host.exec_get_ocr_list() assert result == { diff --git a/tests/unit/agents/test_explorer_prepass.py b/tests/unit/agents/test_explorer_prepass.py index 90c3d74..998a0ad 100644 --- a/tests/unit/agents/test_explorer_prepass.py +++ b/tests/unit/agents/test_explorer_prepass.py @@ -25,7 +25,8 @@ from artemis.agents.explorer.screen_index import ScreenElement, ScreenIndex from artemis.context import ArtemisContext from artemis.graph.state import State -FACADE = "artemis.agents.explorer.explorer" +RUN_SETUP = "artemis.agents.explorer.run_setup" +PERCEPTION_TOOLS = "artemis.agents.explorer.perception_tools" W, H = 1080, 2400 #: One "Settings" button, two "Send" buttons and two OCR-only texts. @@ -108,8 +109,10 @@ def screenshot(tmp_path): @pytest.fixture def no_drawing(): """Keeps the annotation helpers away from the file system.""" + draw_dots = MagicMock() with ( - patch(f"{FACADE}.draw_dots") as draw_dots, + patch(f"{RUN_SETUP}.draw_dots", draw_dots), + patch(f"{PERCEPTION_TOOLS}.draw_dots", draw_dots), patch("pathlib.Path.mkdir"), patch("glob.glob", return_value=[]), ): @@ -127,8 +130,8 @@ async def test_index_is_built_from_state_hierarchy_when_record_is_missing(screen explorer._run_flash = AsyncMock(return_value=_outcome(fallback="nothing")) with ( - patch(f"{FACADE}.StorageManager", return_value=_storage(None)), - patch(f"{FACADE}.perform_ocr", new_callable=AsyncMock) as perform_ocr, + patch(f"{RUN_SETUP}.StorageManager", return_value=_storage(None)), + patch(f"{RUN_SETUP}.perform_ocr", new_callable=AsyncMock) as perform_ocr, ): await explorer.run("gear icon", "", screenshot, _state(screenshot), version="flash") text = (await explorer._search_ui_helper("Settings", prefix="X"))["text"] @@ -145,7 +148,7 @@ async def test_index_ignores_state_hierarchy_for_another_screenshot(screenshot): explorer._run_flash = AsyncMock(return_value=_outcome()) state = _state("/tmp/other.jpg") - with patch(f"{FACADE}.StorageManager", return_value=_storage(None)): + with patch(f"{RUN_SETUP}.StorageManager", return_value=_storage(None)): await explorer.run("Settings", "", screenshot, state, version="flash") assert len(explorer.screen_index) == 0 @@ -224,7 +227,7 @@ async def test_search_by_coords_helper_reports_innermost_first(no_drawing): async def test_get_ocr_list_reads_the_index(no_drawing): explorer = _indexed_explorer() - with patch(f"{FACADE}.StorageManager") as storage_cls: + with patch(f"{RUN_SETUP}.StorageManager") as storage_cls: result = await explorer.exec_get_ocr_list() storage_cls.assert_not_called() @@ -369,7 +372,7 @@ async def test_prepass_answers_unique_exact_match_without_the_engine(screenshot, explorer._run_flash = AsyncMock() explorer._run_loop = AsyncMock() - with patch(f"{FACADE}.StorageManager", return_value=_storage(None)): + with patch(f"{RUN_SETUP}.StorageManager", return_value=_storage(None)): raw = await explorer.run(" settings ", "", screenshot, _state(screenshot), version=version) explorer._run_flash.assert_not_called() @@ -395,7 +398,7 @@ async def test_prepass_leaves_ambiguous_matches_to_the_engine(screenshot): engine = _outcome([{"label": "D1", "coords": [185, 854], "description": "Send"}]) explorer._run_flash = AsyncMock(return_value=engine) - with patch(f"{FACADE}.StorageManager", return_value=_storage(None)): + with patch(f"{RUN_SETUP}.StorageManager", return_value=_storage(None)): raw = await explorer.run("Send", "", screenshot, _state(screenshot), version="flash") explorer._run_flash.assert_awaited_once_with("Send", screenshot) @@ -408,7 +411,7 @@ async def test_prepass_without_pipe_or_match_passes_the_query_through(screenshot engine = _outcome(fallback="Failed to detect: gear icon") explorer._run_flash = AsyncMock(return_value=engine) - with patch(f"{FACADE}.StorageManager", return_value=_storage(None)): + with patch(f"{RUN_SETUP}.StorageManager", return_value=_storage(None)): raw = await explorer.run("gear icon", "", screenshot, _state(screenshot), version="flash") explorer._run_flash.assert_awaited_once_with("gear icon", screenshot) @@ -421,7 +424,7 @@ async def test_prepass_merges_partial_results_with_the_engine_outcome(screenshot detected = {"label": "D1", "coords": [185, 854], "description": "Send"} explorer._run_loop = AsyncMock(return_value=_outcome([detected], fallback="two Send buttons")) - with patch(f"{FACADE}.StorageManager", return_value=_storage(None)): + with patch(f"{RUN_SETUP}.StorageManager", return_value=_storage(None)): raw = await explorer.run( "Settings | Send | Welcome", "", screenshot, _state(screenshot), version="pro" ) @@ -451,9 +454,9 @@ async def test_prepass_never_runs_on_the_fly_ocr_for_flash(screenshot): explorer._run_loop = AsyncMock(return_value=_outcome()) with ( - patch(f"{FACADE}.StorageManager", return_value=_storage(record)), - patch(f"{FACADE}.is_ocr_configured", return_value=True), - patch(f"{FACADE}.perform_ocr", new_callable=AsyncMock, return_value=[]) as perform_ocr, + patch(f"{RUN_SETUP}.StorageManager", return_value=_storage(record)), + patch(f"{RUN_SETUP}.is_ocr_configured", return_value=True), + patch(f"{RUN_SETUP}.perform_ocr", new_callable=AsyncMock, return_value=[]) as perform_ocr, ): await explorer.run("gear icon", "", screenshot, _state(screenshot, []), version="flash") perform_ocr.assert_not_awaited() diff --git a/tests/unit/agents/test_explorer_tiers.py b/tests/unit/agents/test_explorer_tiers.py index 01ac453..e8c400d 100644 --- a/tests/unit/agents/test_explorer_tiers.py +++ b/tests/unit/agents/test_explorer_tiers.py @@ -28,7 +28,9 @@ from artemis.config import settings from artemis.context import ArtemisContext from artemis.graph.state import State -FACADE = "artemis.agents.explorer.explorer" +EXPLORER = "artemis.agents.explorer.explorer" +RUN_SETUP = "artemis.agents.explorer.run_setup" +UNIVERSAL_RUNNER = "artemis.agents.explorer.universal_runner" def _context(model: str = "gemini-3.8-flash", denylist: list[str] | None = None) -> MagicMock: @@ -108,7 +110,7 @@ def _isolated_env(monkeypatch): def test_construction_creates_no_genai_client(): - with patch(f"{FACADE}.genai.Client") as client_cls: + with patch(f"{EXPLORER}.genai.Client") as client_cls: explorer = Explorer(_context()) client_cls.assert_not_called() assert explorer.client is None @@ -124,7 +126,7 @@ def test_construction_creates_no_genai_client(): ) def test_tier_sets_turn_budget_and_exposed_tools(version, max_turns, expected_tools): explorer = Explorer(_context()) - with patch(f"{FACADE}.is_ocr_configured", return_value=True): + with patch(f"{EXPLORER}.is_ocr_configured", return_value=True): explorer._apply_tier(get_tier(version)) native_names = {t.name for t in explorer.get_exposed_tools()} universal_names = {t["function"]["name"] for t in explorer._universal_exposed_tools()} @@ -137,7 +139,7 @@ def test_tier_sets_turn_budget_and_exposed_tools(version, max_turns, expected_to def test_unconfigured_ocr_and_user_denylist_hide_tools_on_ultra(): explorer = Explorer(_context(denylist=["inspect_region"])) - with patch(f"{FACADE}.is_ocr_configured", return_value=False): + with patch(f"{EXPLORER}.is_ocr_configured", return_value=False): explorer._apply_tier(get_tier("ultra")) names = {t.name for t in explorer.get_exposed_tools()} assert "get_ocr_list" not in names @@ -190,9 +192,9 @@ async def test_universal_run_without_google_key_never_touches_genai(screenshot): with ( patch.object(settings, "GOOGLE_API_KEY", None), - patch(f"{FACADE}.genai.Client") as client_cls, - patch(f"{FACADE}.get_llm", return_value=llm), - patch(f"{FACADE}.StorageManager", return_value=storage), + patch(f"{EXPLORER}.genai.Client") as client_cls, + patch(f"{UNIVERSAL_RUNNER}.get_llm", return_value=llm), + patch(f"{RUN_SETUP}.StorageManager", return_value=storage), ): explorer = Explorer(ctx) raw = await explorer.run("find it", "", screenshot, _state(screenshot), version="pro") @@ -342,7 +344,7 @@ def test_pro_prompt_only_describes_exposed_tools(): def test_ultra_prompt_describes_all_tools(): explorer = Explorer(_context()) - with patch(f"{FACADE}.is_ocr_configured", return_value=True): + with patch(f"{EXPLORER}.is_ocr_configured", return_value=True): explorer._apply_tier(get_tier("ultra")) prompt, error = explorer._build_prompt_template(get_tier("ultra")) assert error is None @@ -355,7 +357,7 @@ def test_ultra_prompt_describes_all_tools(): def test_prompt_denylist_section_only_for_user_denied_tier_tools(): explorer = Explorer(_context(denylist=["inspect_region", "search_xml_ocr"])) - with patch(f"{FACADE}.is_ocr_configured", return_value=True): + with patch(f"{EXPLORER}.is_ocr_configured", return_value=True): explorer._apply_tier(get_tier("ultra")) ultra_prompt, _ = explorer._build_prompt_template(get_tier("ultra")) explorer._apply_tier(get_tier("pro")) @@ -399,8 +401,8 @@ async def test_flash_run_returns_detected_candidates(screenshot): } ) with ( - patch(f"{FACADE}._run_object_detection", detection), - patch(f"{FACADE}.genai.Client") as client_cls, + patch(f"{RUN_SETUP}._run_object_detection", detection), + patch(f"{EXPLORER}.genai.Client") as client_cls, ): explorer = Explorer(_context()) raw = await explorer.run( @@ -421,7 +423,7 @@ async def test_flash_run_returns_detected_candidates(screenshot): @pytest.mark.asyncio async def test_flash_run_reports_missing_detection(screenshot): - with patch(f"{FACADE}._run_object_detection", AsyncMock(return_value={"detected": []})): + with patch(f"{RUN_SETUP}._run_object_detection", AsyncMock(return_value={"detected": []})): raw = await Explorer(_context()).run( "ghost", "", screenshot, _state(screenshot), version="flash" ) @@ -442,8 +444,8 @@ async def test_native_ultra_run_uses_async_cache_api(screenshot): storage.get_image.return_value = None with ( - patch(f"{FACADE}.genai.Client") as client_cls, - patch(f"{FACADE}.StorageManager", return_value=storage), + patch(f"{EXPLORER}.genai.Client") as client_cls, + patch(f"{RUN_SETUP}.StorageManager", return_value=storage), patch.object(settings, "EXPLORER_CACHING", None), ): explorer = Explorer(ctx) @@ -466,7 +468,7 @@ async def test_native_pro_run_skips_cache_by_default(screenshot): storage.get_image.return_value = None with ( - patch(f"{FACADE}.StorageManager", return_value=storage), + patch(f"{RUN_SETUP}.StorageManager", return_value=storage), patch.object(settings, "EXPLORER_CACHING", None), ): await Explorer(ctx).run("target", "", screenshot, _state(screenshot), version="pro") @@ -489,7 +491,7 @@ async def test_native_file_api_uploads_and_deletes_asynchronously(screenshot, mo storage = MagicMock() storage.get_image.return_value = None - with patch(f"{FACADE}.StorageManager", return_value=storage): + with patch(f"{RUN_SETUP}.StorageManager", return_value=storage): await Explorer(ctx).run("target", "", screenshot, _state(screenshot), version="pro") client.aio.files.upload.assert_awaited_once_with(file=screenshot) @@ -507,7 +509,7 @@ async def test_run_closes_http_client_on_both_engines(screenshot): ctx._genai_client = client native = Explorer(ctx) native.http_client = SimpleNamespace(aclose=AsyncMock()) - with patch(f"{FACADE}.StorageManager", return_value=storage): + with patch(f"{RUN_SETUP}.StorageManager", return_value=storage): await native.run("t", "", screenshot, _state(screenshot), version="pro") native_close = native.http_client # already reset assert native_close is None @@ -519,8 +521,8 @@ async def test_run_closes_http_client_on_both_engines(screenshot): http_client = SimpleNamespace(aclose=AsyncMock()) universal.http_client = http_client with ( - patch(f"{FACADE}.StorageManager", return_value=storage), - patch(f"{FACADE}.get_llm", return_value=llm), + patch(f"{RUN_SETUP}.StorageManager", return_value=storage), + patch(f"{UNIVERSAL_RUNNER}.get_llm", return_value=llm), pytest.raises(RuntimeError), ): await universal.run("t", "", screenshot, _state(screenshot), version="pro") diff --git a/tests/unit/agents/test_explorer_universal.py b/tests/unit/agents/test_explorer_universal.py index 5b7f811..65e86e4 100644 --- a/tests/unit/agents/test_explorer_universal.py +++ b/tests/unit/agents/test_explorer_universal.py @@ -20,10 +20,8 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest from langchain_core.messages import AIMessage -from artemis.agents.explorer.explorer import ( - UNIVERSAL_EXPLORER_TOOLS, - Explorer, -) +from artemis.agents.explorer.explorer import Explorer +from artemis.agents.explorer.tool_declarations import UNIVERSAL_EXPLORER_TOOLS from artemis.config import settings from artemis.context import ArtemisContext from artemis.graph.state import State @@ -106,7 +104,7 @@ async def test_run_universal_explorer_success(mock_context, tmp_path): mock_llm.bind_tools.return_value = mock_bound with patch( - "artemis.agents.explorer.explorer.get_llm", + "artemis.agents.explorer.universal_runner.get_llm", return_value=mock_llm, ): mock_state = Mock(spec=State) diff --git a/tests/unit/agents/test_operator_prompts.py b/tests/unit/agents/test_operator_prompts.py index 19a2727..06b5346 100644 --- a/tests/unit/agents/test_operator_prompts.py +++ b/tests/unit/agents/test_operator_prompts.py @@ -12,11 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -from artemis.agents.operator.prompt_builder import ( - OperatorPromptBuilder, +from artemis.agents.operator.prompts import ( + apply_operator_prompt_contract, load_operator_prompts, ) -from artemis.agents.operator.prompts import apply_operator_prompt_contract def test_operator_prompt_contract_matches_note_runtime_semantics(): @@ -68,9 +67,8 @@ def test_operator_prompt_keeps_large_list_traversal_single_pass_until_boundary() assert "Terminate exploration when a definitive boundary is reached" not in prompt -def test_compatibility_builder_uses_canonical_operator_template(): +def test_contract_strips_legacy_tool_literals(): expected = apply_operator_prompt_contract(load_operator_prompts()["main_template"]) - assert OperatorPromptBuilder.build_system_message() == expected assert "wait_for_delay(seconds=" not in expected assert 'press_key(key="home"' not in expected diff --git a/tests/unit/agents/test_operator_transcript.py b/tests/unit/agents/test_operator_transcript.py index 2966dae..ca1e754 100644 --- a/tests/unit/agents/test_operator_transcript.py +++ b/tests/unit/agents/test_operator_transcript.py @@ -33,7 +33,7 @@ import pytest from langchain_core.messages import HumanMessage, SystemMessage from artemis.agents.operator.operator import OperatorNode -from artemis.agents.operator.prompt_builder import load_operator_prompts +from artemis.agents.operator.prompts import load_operator_prompts from artemis.agents.operator.prompts import ( PLAN_HISTORY_STATIC_POINTER, PLAN_HISTORY_TEMPLATE_SECTION, diff --git a/tests/unit/agents/test_prompt_assembly.py b/tests/unit/agents/test_prompt_assembly.py index e9d05e4..04ce9b9 100644 --- a/tests/unit/agents/test_prompt_assembly.py +++ b/tests/unit/agents/test_prompt_assembly.py @@ -20,7 +20,7 @@ import re from jinja2 import Template import pytest -from artemis.agents.operator.prompt_builder import load_operator_prompts +from artemis.agents.operator.prompts import load_operator_prompts from artemis.agents.operator.prompts import ( OPERATOR_PROMPT_TOOLSET, apply_operator_prompt_contract, diff --git a/tests/unit/mcp/test_action_specs.py b/tests/unit/mcp/test_action_specs.py index cbdca86..3bed8bf 100644 --- a/tests/unit/mcp/test_action_specs.py +++ b/tests/unit/mcp/test_action_specs.py @@ -208,20 +208,14 @@ def test_operator_verbs_lower_onto_manifest_actions(): ) -def test_registry_wait_for_delay_seconds_drift_stays_dead(): - """Guards against re-growing a shadow device-action registration. +def test_manifest_wait_for_delay_uses_milliseconds(): + """The manifest's ``wait_for_delay`` takes ``time_in_ms``, never *seconds*. ``artemis/tools/actions/device_actions.py`` used to shadow-register a second - ``wait_for_delay`` taking *seconds* while every prompt teaches milliseconds - (the last such surface, ``artemis/tools/mobile/exec_tools.py``, is deleted). - The ToolRegistry must never again carry a ``wait_for_delay`` whose unit - disagrees with the manifest's ``time_in_ms``. + ``wait_for_delay`` taking *seconds* while every prompt teaches milliseconds. + The shadow ToolRegistry channel is deleted; the manifest is the only surface + left, and its unit must stay milliseconds. """ - from artemis.tools.base import ToolRegistry - - tool = ToolRegistry.get("wait_for_delay") - if tool is None: - return # Not registered at all: nothing to drift. - fields = set(tool.args_schema.model_json_schema().get("properties", {})) + fields = set(tool_declaration("wait_for_delay").parameters["properties"]) assert "seconds" not in fields assert "time_in_ms" in fields diff --git a/tests/unit/runtime/test_awake_service.py b/tests/unit/runtime/test_awake_service.py index 1c29043..9956f5d 100644 --- a/tests/unit/runtime/test_awake_service.py +++ b/tests/unit/runtime/test_awake_service.py @@ -29,7 +29,10 @@ def test_usb_policy_is_primary_when_android_reports_it_active(mock_run, lease_ty lease_type.assert_called_once_with("device-123") lease_type.return_value.cleanup_unowned_references.assert_called_once_with() commands = [call.args[0] for call in mock_run.call_args_list] - assert ["shell", "svc", "power", "stayon", "usb"] in [command[3:] for command in commands] + # Commands are endpoint-qualified (adb -H host -P port -s serial ...); + # assert on the payload after the serial rather than a fixed prefix width. + payloads = [command[command.index("-s") + 2 :] for command in commands if "-s" in command] + assert ["shell", "svc", "power", "stayon", "usb"] in payloads assert not any("SCREEN_BRIGHT_WAKE_LOCK" in command for command in commands) assert not any("KEYCODE_UNKNOWN" in command for command in commands) diff --git a/tests/unit/test_prompt_builder.py b/tests/unit/test_prompt_builder.py deleted file mode 100644 index 40773ef..0000000 --- a/tests/unit/test_prompt_builder.py +++ /dev/null @@ -1,34 +0,0 @@ -# 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. - -"""Unit tests for OperatorPromptBuilder.""" - -from artemis.agents.operator.prompt_builder import OperatorPromptBuilder - - -def test_operator_prompt_builder(): - """Verify prompt rendering formats variables correctly.""" - system_msg = OperatorPromptBuilder.build_system_message() - assert "Android Device Operator" in system_msg - assert "click" in system_msg - - human_msg = OperatorPromptBuilder.build_human_message( - goal="Open Settings", - sub_goal="Tap Network & internet", - current_turn=2, - history="1. Tapped Settings icon", - ) - assert "Open Settings" in human_msg - assert "Tap Network & internet" in human_msg - assert "Current Turn: 2" in human_msg diff --git a/tests/unit/tools/test_append_note_pure_tool.py b/tests/unit/tools/test_append_note_pure_tool.py index 428dc77..373285e 100644 --- a/tests/unit/tools/test_append_note_pure_tool.py +++ b/tests/unit/tools/test_append_note_pure_tool.py @@ -16,7 +16,7 @@ from pathlib import Path from unittest.mock import MagicMock, patch from artemis.context import ArtemisContext -from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.base import ArtemisTool from artemis.tools.scratchpad import ( AppendNoteArgs, AppendNotePure, @@ -35,8 +35,8 @@ def mock_ctx(tmp_path): return ctx -def test_append_note_pure_tool_subclass_and_registry(): - """Verify AppendNotePureTool is a subclass of ArtemisTool and properly registered.""" +def test_append_note_pure_tool_subclass(): + """Verify AppendNotePureTool is a subclass of ArtemisTool.""" assert issubclass(AppendNotePureTool, ArtemisTool) assert issubclass(AppendNotePure, ArtemisTool) assert isinstance(append_note_pure, ArtemisTool) @@ -46,11 +46,6 @@ def test_append_note_pure_tool_subclass_and_registry(): assert append_note_pure.category == "memory" assert append_note_pure.args_schema == AppendNoteArgs - # Registry lookup - reg_tool = ToolRegistry.get("append_note_pure") - assert reg_tool is not None - assert isinstance(reg_tool, AppendNotePureTool) - # GenAI FunctionDeclaration export declaration = append_note_pure.to_genai_declaration() assert declaration.name == "append_note_pure" diff --git a/tests/unit/tools/test_append_note_tool.py b/tests/unit/tools/test_append_note_tool.py index bf618c1..83c9ab1 100644 --- a/tests/unit/tools/test_append_note_tool.py +++ b/tests/unit/tools/test_append_note_tool.py @@ -17,7 +17,7 @@ from unittest.mock import MagicMock, patch from artemis.context import ArtemisContext from artemis.graph.state import State -from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.base import ArtemisTool from artemis.tools.scratchpad import ( AppendNote, AppendNoteArgs, @@ -39,8 +39,8 @@ def mock_ctx(tmp_path): return ctx -def test_append_note_tool_subclass_and_registry(): - """Verify AppendNoteTool is a subclass of ArtemisTool and properly registered.""" +def test_append_note_tool_subclass(): + """Verify AppendNoteTool is a subclass of ArtemisTool.""" assert issubclass(AppendNoteTool, ArtemisTool) assert issubclass(AppendNote, ArtemisTool) assert isinstance(append_note, ArtemisTool) @@ -50,11 +50,6 @@ def test_append_note_tool_subclass_and_registry(): assert append_note.category == "memory" assert append_note.args_schema == AppendNoteArgs - # Registry lookup - reg_tool = ToolRegistry.get("append_note") - assert reg_tool is not None - assert isinstance(reg_tool, AppendNoteTool) - # GenAI FunctionDeclaration export declaration = append_note.to_genai_declaration() assert declaration.name == "append_note" diff --git a/tests/unit/tools/test_command_tool.py b/tests/unit/tools/test_command_tool.py index bbef89e..4dfa915 100644 --- a/tests/unit/tools/test_command_tool.py +++ b/tests/unit/tools/test_command_tool.py @@ -373,9 +373,9 @@ async def test_analyze_task_output_tool(mock_get_llm, mock_ctx): del _FINISHED_TASKS_LOGS[task_id] -def test_run_adb_command_tool_subclass_and_registry(): - """Verify RunAdbCommandTool is a subclass of ArtemisTool and properly registered.""" - from artemis.tools.base import ArtemisTool, ToolRegistry +def test_run_adb_command_tool_subclass(): + """Verify RunAdbCommandTool is a subclass of ArtemisTool.""" + from artemis.tools.base import ArtemisTool from artemis.tools.command_tool import ( RunAdbCommand, RunAdbCommandArgs, @@ -392,11 +392,6 @@ def test_run_adb_command_tool_subclass_and_registry(): assert run_adb_command.category == "system" assert run_adb_command.args_schema == RunAdbCommandArgs - # Registry lookup - reg_tool = ToolRegistry.get("run_adb_command") - assert reg_tool is not None - assert isinstance(reg_tool, RunAdbCommandTool) - # GenAI FunctionDeclaration export declaration = run_adb_command.to_genai_declaration() assert declaration.name == "run_adb_command" @@ -421,9 +416,9 @@ async def test_run_adb_command_tool_mock_driver_execute(): assert "pm list packages" in driver.action_history[-1]["command"] -def test_manage_task_tool_subclass_and_registry(): - """Verify ManageTaskTool is a subclass of ArtemisTool and properly registered.""" - from artemis.tools.base import ArtemisTool, ToolRegistry +def test_manage_task_tool_subclass(): + """Verify ManageTaskTool is a subclass of ArtemisTool.""" + from artemis.tools.base import ArtemisTool from artemis.tools.command_tool import ( ManageTask, ManageTaskArgs, @@ -440,11 +435,6 @@ def test_manage_task_tool_subclass_and_registry(): assert manage_task.category == "system" assert manage_task.args_schema == ManageTaskArgs - # Registry lookup - reg_tool = ToolRegistry.get("manage_task") - assert reg_tool is not None - assert isinstance(reg_tool, ManageTaskTool) - # GenAI FunctionDeclaration export declaration = manage_task.to_genai_declaration() assert declaration.name == "manage_task" @@ -488,9 +478,9 @@ async def test_manage_task_tool_direct_execute(mock_ctx): registry.background.clear() -def test_analyze_task_output_tool_subclass_and_registry(): - """Verify AnalyzeTaskOutputTool is a subclass of ArtemisTool and properly registered.""" - from artemis.tools.base import ArtemisTool, ToolRegistry +def test_analyze_task_output_tool_subclass(): + """Verify AnalyzeTaskOutputTool is a subclass of ArtemisTool.""" + from artemis.tools.base import ArtemisTool from artemis.tools.command_tool import ( AnalyzeTaskOutput, AnalyzeTaskOutputArgs, @@ -507,11 +497,6 @@ def test_analyze_task_output_tool_subclass_and_registry(): assert analyze_task_output.category == "system" assert analyze_task_output.args_schema == AnalyzeTaskOutputArgs - # Registry lookup - reg_tool = ToolRegistry.get("analyze_task_output") - assert reg_tool is not None - assert isinstance(reg_tool, AnalyzeTaskOutputTool) - # GenAI FunctionDeclaration export declaration = analyze_task_output.to_genai_declaration() assert declaration.name == "analyze_task_output" @@ -558,9 +543,9 @@ async def test_analyze_task_output_tool_direct_execute(mock_get_llm, mock_ctx): del _FINISHED_TASKS_LOGS[task_id] -def test_run_short_adb_command_tool_subclass_and_registry(): - """Verify RunShortAdbCommandTool is a subclass of ArtemisTool and properly registered.""" - from artemis.tools.base import ArtemisTool, ToolRegistry +def test_run_short_adb_command_tool_subclass(): + """Verify RunShortAdbCommandTool is a subclass of ArtemisTool.""" + from artemis.tools.base import ArtemisTool from artemis.tools.command_tool import ( RunShortAdbCommand, RunShortAdbCommandArgs, @@ -577,11 +562,6 @@ def test_run_short_adb_command_tool_subclass_and_registry(): assert run_short_adb_command.category == "system" assert run_short_adb_command.args_schema == RunShortAdbCommandArgs - # Registry lookup - reg_tool = ToolRegistry.get("run_short_adb_command") - assert reg_tool is not None - assert isinstance(reg_tool, RunShortAdbCommandTool) - # GenAI FunctionDeclaration export declaration = run_short_adb_command.to_genai_declaration() assert declaration.name == "run_short_adb_command" diff --git a/tests/unit/tools/test_committee_tool.py b/tests/unit/tools/test_committee_tool.py index 7e6163e..27e3d5e 100644 --- a/tests/unit/tools/test_committee_tool.py +++ b/tests/unit/tools/test_committee_tool.py @@ -168,9 +168,9 @@ class TestCommitteeTool(unittest.IsolatedAsyncioTestCase): # Clean up temp directory shutil.rmtree(temp_dir) - def test_ask_committee_tool_subclass_and_registry(self): - """Verify AskCommitteeTool is a subclass of ArtemisTool and properly registered.""" - from artemis.tools.base import ArtemisTool, ToolRegistry + def test_ask_committee_tool_subclass(self): + """Verify AskCommitteeTool is a subclass of ArtemisTool.""" + from artemis.tools.base import ArtemisTool from artemis.tools.committee_tool import ( AskCommittee, AskCommitteeArgs, @@ -191,11 +191,6 @@ class TestCommitteeTool(unittest.IsolatedAsyncioTestCase): self.assertEqual(ask_committee.category, "custom") self.assertEqual(ask_committee.args_schema, AskCommitteeArgs) - # Registry lookup - reg_tool = ToolRegistry.get("ask_committee") - self.assertIsNotNone(reg_tool) - self.assertIsInstance(reg_tool, AskCommitteeTool) - # GenAI FunctionDeclaration export declaration = ask_committee.to_genai_declaration() self.assertEqual(declaration.name, "ask_committee") diff --git a/tests/unit/tools/test_explorer_tool.py b/tests/unit/tools/test_explorer_tool.py index e6995d3..265e2b0 100644 --- a/tests/unit/tools/test_explorer_tool.py +++ b/tests/unit/tools/test_explorer_tool.py @@ -27,7 +27,7 @@ from artemis.agents.explorer.constants import ( ) from artemis.context import ArtemisContext from artemis.graph.state import State -from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.base import ArtemisTool from artemis.tools.explorer_tool import ( AskExplorerArgs, AskExplorerTool, @@ -122,9 +122,7 @@ def test_description_does_not_depend_on_the_tier(version): assert tool.version == version -def test_registry_and_wrapper_export(): - reg_tool = ToolRegistry.get("ask_explorer") - assert isinstance(reg_tool, AskExplorerTool) +def test_wrapper_export(): assert ask_explorer_wrapper.tool_fn_getter is get_ask_explorer_tool assert ask_explorer_wrapper.on_success_fn("x") == "x" assert "boom" in ask_explorer_wrapper.on_failure_fn("boom") diff --git a/tests/unit/tools/test_image_processor_tool.py b/tests/unit/tools/test_image_processor_tool.py index 03b8072..8bc200f 100644 --- a/tests/unit/tools/test_image_processor_tool.py +++ b/tests/unit/tools/test_image_processor_tool.py @@ -64,9 +64,9 @@ async def test_ask_vision_coder_success(): assert result == vision_coder_outcome -def test_ask_image_processor_tool_subclass_and_registry(): - """Verify AskImageProcessorTool is an ArtemisTool subclass and registered in ToolRegistry.""" - from artemis.tools.base import ArtemisTool, ToolRegistry +def test_ask_image_processor_tool_subclass(): + """Verify AskImageProcessorTool is an ArtemisTool subclass.""" + from artemis.tools.base import ArtemisTool from artemis.tools.image_processor_tool import ( AskImageProcessor, AskImageProcessorTool, @@ -87,11 +87,6 @@ def test_ask_image_processor_tool_subclass_and_registry(): assert ask_image_processor.category == "custom" assert ask_image_processor.args_schema == AskVisionCoderArgs - # Registry lookup - reg_tool = ToolRegistry.get("ask_image_processor") - assert reg_tool is not None - assert isinstance(reg_tool, AskImageProcessorTool) - # GenAI FunctionDeclaration export declaration = ask_image_processor.to_genai_declaration() assert declaration.name == "ask_image_processor" diff --git a/tests/unit/tools/test_launch_app_tool.py b/tests/unit/tools/test_launch_app_tool.py index 4c21443..0c1a107 100644 --- a/tests/unit/tools/test_launch_app_tool.py +++ b/tests/unit/tools/test_launch_app_tool.py @@ -17,7 +17,7 @@ from unittest.mock import AsyncMock, MagicMock, patch from artemis.context import ArtemisContext from artemis.drivers.base import BaseDeviceDriver from artemis.graph.state import State -from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.base import ArtemisTool from artemis.tools.mobile.launch_app import ( LaunchApp, LaunchAppArgs, @@ -49,8 +49,8 @@ def mock_state(): return MagicMock(spec=State) -def test_launch_app_tool_subclass_and_registry(): - """Verify LaunchAppTool is a subclass of ArtemisTool and properly registered.""" +def test_launch_app_tool_subclass(): + """Verify LaunchAppTool is a subclass of ArtemisTool.""" assert issubclass(LaunchAppTool, ArtemisTool) assert issubclass(LaunchApp, ArtemisTool) assert isinstance(launch_app, ArtemisTool) @@ -60,11 +60,6 @@ def test_launch_app_tool_subclass_and_registry(): assert launch_app.category == "action" assert launch_app.args_schema == LaunchAppArgs - # Registry lookup - reg_tool = ToolRegistry.get("launch_app") - assert reg_tool is not None - assert isinstance(reg_tool, LaunchAppTool) - # GenAI FunctionDeclaration export declaration = launch_app.to_genai_declaration() assert declaration.name == "launch_app" diff --git a/tests/unit/tools/test_list_notes_pure_tool.py b/tests/unit/tools/test_list_notes_pure_tool.py index 997115e..79d2555 100644 --- a/tests/unit/tools/test_list_notes_pure_tool.py +++ b/tests/unit/tools/test_list_notes_pure_tool.py @@ -15,7 +15,7 @@ from unittest.mock import MagicMock, patch from artemis.context import ArtemisContext -from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.base import ArtemisTool from artemis.tools.scratchpad import ( ListNotesArgs, ListNotesPure, @@ -35,8 +35,8 @@ def mock_ctx(tmp_path): return ctx -def test_list_notes_pure_tool_subclass_and_registry(): - """Verify ListNotesPureTool is a subclass of ArtemisTool and properly registered.""" +def test_list_notes_pure_tool_subclass(): + """Verify ListNotesPureTool is a subclass of ArtemisTool.""" assert issubclass(ListNotesPureTool, ArtemisTool) assert issubclass(ListNotesPure, ArtemisTool) assert isinstance(list_notes_pure, ArtemisTool) @@ -46,11 +46,6 @@ def test_list_notes_pure_tool_subclass_and_registry(): assert list_notes_pure.category == "memory" assert list_notes_pure.args_schema == ListNotesArgs - # Registry lookup - reg_tool = ToolRegistry.get("list_notes_pure") - assert reg_tool is not None - assert isinstance(reg_tool, ListNotesPureTool) - # GenAI FunctionDeclaration export declaration = list_notes_pure.to_genai_declaration() assert declaration.name == "list_notes_pure" diff --git a/tests/unit/tools/test_list_notes_tool.py b/tests/unit/tools/test_list_notes_tool.py index 224cf3d..1a8ab6a 100644 --- a/tests/unit/tools/test_list_notes_tool.py +++ b/tests/unit/tools/test_list_notes_tool.py @@ -16,7 +16,7 @@ from unittest.mock import MagicMock, patch from artemis.context import ArtemisContext from artemis.graph.state import State -from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.base import ArtemisTool from artemis.tools.scratchpad import ( ListNotes, ListNotesArgs, @@ -39,8 +39,8 @@ def mock_ctx(tmp_path): return ctx -def test_list_notes_tool_subclass_and_registry(): - """Verify ListNotesTool is a subclass of ArtemisTool and properly registered.""" +def test_list_notes_tool_subclass(): + """Verify ListNotesTool is a subclass of ArtemisTool.""" assert issubclass(ListNotesTool, ArtemisTool) assert issubclass(ListNotes, ArtemisTool) assert isinstance(list_notes, ArtemisTool) @@ -50,11 +50,6 @@ def test_list_notes_tool_subclass_and_registry(): assert list_notes.category == "memory" assert list_notes.args_schema == ListNotesArgs - # Registry lookup - reg_tool = ToolRegistry.get("list_notes") - assert reg_tool is not None - assert isinstance(reg_tool, ListNotesTool) - # GenAI FunctionDeclaration export declaration = list_notes.to_genai_declaration() assert declaration.name == "list_notes" diff --git a/tests/unit/tools/test_log_tool.py b/tests/unit/tools/test_log_tool.py index 990d7bd..98896d8 100644 --- a/tests/unit/tools/test_log_tool.py +++ b/tests/unit/tools/test_log_tool.py @@ -16,7 +16,7 @@ from unittest.mock import AsyncMock, MagicMock, patch from artemis.context import ArtemisContext from artemis.graph.state import State -from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.base import ArtemisTool from artemis.tools.log_tool import ( AnalyzeLogs, AnalyzeLogsArgs, @@ -88,8 +88,8 @@ async def test_analyze_logs_tool_no_context(): assert result == "Error: ArtemisContext is required for LogAnalyzer." -def test_analyze_logs_tool_subclass_and_registry(): - """Verify AnalyzeLogsTool is an ArtemisTool subclass and registered in ToolRegistry.""" +def test_analyze_logs_tool_subclass(): + """Verify AnalyzeLogsTool is an ArtemisTool subclass.""" assert issubclass(AnalyzeLogsTool, ArtemisTool) assert issubclass(AnalyzeLogs, ArtemisTool) assert issubclass(LogTool, ArtemisTool) @@ -100,11 +100,6 @@ def test_analyze_logs_tool_subclass_and_registry(): assert analyze_logs.category == "custom" assert analyze_logs.args_schema == AnalyzeLogsArgs - # Registry lookup - reg_tool = ToolRegistry.get("analyze_logs") - assert reg_tool is not None - assert isinstance(reg_tool, AnalyzeLogsTool) - # GenAI FunctionDeclaration export declaration = analyze_logs.to_genai_declaration() assert declaration.name == "analyze_logs" diff --git a/tests/unit/tools/test_object_detection_tool.py b/tests/unit/tools/test_object_detection_tool.py index ffd0f35..57b6fa3 100644 --- a/tests/unit/tools/test_object_detection_tool.py +++ b/tests/unit/tools/test_object_detection_tool.py @@ -21,7 +21,7 @@ from langchain_core.messages import ToolMessage from artemis.context import ArtemisContext from artemis.graph.state import State -from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.base import ArtemisTool from artemis.tools.object_detection_tool import ( ObjectDetection, ObjectDetectionArgs, @@ -41,7 +41,7 @@ from artemis.tools.object_detection_tool import ( import pytest -def test_object_detection_tool_subclass_and_registry(): +def test_object_detection_tool_subclass(): """Verify ObjectDetectionTool and OperatorObjectDetectionTool subclass ArtemisTool and register.""" assert issubclass(ObjectDetectionTool, ArtemisTool) assert issubclass(ObjectDetection, ArtemisTool) @@ -65,11 +65,6 @@ def test_object_detection_tool_subclass_and_registry(): assert operator_object_detection.category == "perception" assert operator_object_detection.args_schema == OperatorObjectDetectionArgs - # Registry lookup - reg_tool = ToolRegistry.get("object_detection") - assert reg_tool is not None - assert isinstance(reg_tool, ObjectDetectionTool) - # GenAI FunctionDeclaration export declaration = object_detection.to_genai_declaration() assert declaration.name == "object_detection" diff --git a/tests/unit/tools/test_ocr_tool.py b/tests/unit/tools/test_ocr_tool.py index 9eb67a6..720ab26 100644 --- a/tests/unit/tools/test_ocr_tool.py +++ b/tests/unit/tools/test_ocr_tool.py @@ -18,7 +18,7 @@ from unittest.mock import AsyncMock, MagicMock, patch from artemis.context import ArtemisContext from artemis.drivers.base import BaseDeviceDriver from artemis.graph.state import State -from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.base import ArtemisTool from artemis.tools.mobile.ocr import ( OCRRecognition, OcrArgs, @@ -63,8 +63,8 @@ def mock_state(tmp_path): return state -def test_ocr_tool_subclass_and_registry(): - """Verify OcrRecognitionTool is a subclass of ArtemisTool and properly registered.""" +def test_ocr_tool_subclass(): + """Verify OcrRecognitionTool is a subclass of ArtemisTool.""" assert issubclass(OcrRecognitionTool, ArtemisTool) assert issubclass(OcrRecognition, ArtemisTool) assert issubclass(OCRRecognition, ArtemisTool) @@ -76,11 +76,6 @@ def test_ocr_tool_subclass_and_registry(): assert ocr_recognition.category == "explorer" assert ocr_recognition.args_schema == OcrArgs - # Registry lookup - reg_tool = ToolRegistry.get("ocr_recognition") - assert reg_tool is not None - assert isinstance(reg_tool, OcrRecognitionTool) - # GenAI FunctionDeclaration export declaration = ocr_recognition.to_genai_declaration() assert declaration.name == "ocr_recognition" diff --git a/tests/unit/tools/test_read_hierarchy_tool.py b/tests/unit/tools/test_read_hierarchy_tool.py index 9ee542a..03e639e 100644 --- a/tests/unit/tools/test_read_hierarchy_tool.py +++ b/tests/unit/tools/test_read_hierarchy_tool.py @@ -16,7 +16,7 @@ from unittest.mock import AsyncMock, MagicMock from artemis.context import ArtemisContext from artemis.drivers.base import BaseDeviceDriver -from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.base import ArtemisTool from artemis.tools.mobile.read_hierarchy import ( GetUIHierarchy, GetUiHierarchy, @@ -46,8 +46,8 @@ def mock_driver(): return driver -def test_get_ui_hierarchy_tool_subclass_and_registry(): - """Verify GetUiHierarchyTool is a subclass of ArtemisTool and properly registered.""" +def test_get_ui_hierarchy_tool_subclass(): + """Verify GetUiHierarchyTool is a subclass of ArtemisTool.""" assert issubclass(GetUiHierarchyTool, ArtemisTool) assert issubclass(GetUiHierarchy, ArtemisTool) assert issubclass(GetUIHierarchy, ArtemisTool) @@ -58,11 +58,6 @@ def test_get_ui_hierarchy_tool_subclass_and_registry(): assert get_ui_hierarchy.category == "diagnostic" assert get_ui_hierarchy.args_schema == GetUiHierarchyArgs - # Registry lookup - reg_tool = ToolRegistry.get("get_ui_hierarchy") - assert reg_tool is not None - assert isinstance(reg_tool, GetUiHierarchyTool) - # GenAI FunctionDeclaration export declaration = get_ui_hierarchy.to_genai_declaration() assert declaration.name == "get_ui_hierarchy" diff --git a/tests/unit/tools/test_read_logs_tool.py b/tests/unit/tools/test_read_logs_tool.py index 5e14340..b3ac417 100644 --- a/tests/unit/tools/test_read_logs_tool.py +++ b/tests/unit/tools/test_read_logs_tool.py @@ -15,7 +15,7 @@ from unittest.mock import MagicMock, patch from artemis.context import ArtemisContext -from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.base import ArtemisTool from artemis.tools.mobile.read_logs import ( ReadLogs, ReadLogsArgs, @@ -33,8 +33,8 @@ def mock_ctx(): return ctx -def test_read_logs_tool_subclass_and_registry(): - """Verify ReadLogsTool is a subclass of ArtemisTool and properly registered.""" +def test_read_logs_tool_subclass(): + """Verify ReadLogsTool is a subclass of ArtemisTool.""" assert issubclass(ReadLogsTool, ArtemisTool) assert issubclass(ReadLogs, ArtemisTool) assert issubclass(ReadLogsToolAlias, ArtemisTool) @@ -45,11 +45,6 @@ def test_read_logs_tool_subclass_and_registry(): assert read_logs.category == "diagnostic" assert read_logs.args_schema == ReadLogsArgs - # Registry lookup - reg_tool = ToolRegistry.get("read_logs") - assert reg_tool is not None - assert isinstance(reg_tool, ReadLogsTool) - # GenAI FunctionDeclaration export declaration = read_logs.to_genai_declaration() assert declaration.name == "read_logs" diff --git a/tests/unit/tools/test_read_note_pure_tool.py b/tests/unit/tools/test_read_note_pure_tool.py index f8bb661..22d2fde 100644 --- a/tests/unit/tools/test_read_note_pure_tool.py +++ b/tests/unit/tools/test_read_note_pure_tool.py @@ -15,7 +15,7 @@ from unittest.mock import MagicMock, patch from artemis.context import ArtemisContext -from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.base import ArtemisTool from artemis.tools.scratchpad import ( ReadNoteArgs, ReadNotePure, @@ -35,8 +35,8 @@ def mock_ctx(tmp_path): return ctx -def test_read_note_pure_tool_subclass_and_registry(): - """Verify ReadNotePureTool is a subclass of ArtemisTool and properly registered.""" +def test_read_note_pure_tool_subclass(): + """Verify ReadNotePureTool is a subclass of ArtemisTool.""" assert issubclass(ReadNotePureTool, ArtemisTool) assert issubclass(ReadNotePure, ArtemisTool) assert isinstance(read_note_pure, ArtemisTool) @@ -46,11 +46,6 @@ def test_read_note_pure_tool_subclass_and_registry(): assert read_note_pure.category == "memory" assert read_note_pure.args_schema == ReadNoteArgs - # Registry lookup - reg_tool = ToolRegistry.get("read_note_pure") - assert reg_tool is not None - assert isinstance(reg_tool, ReadNotePureTool) - # GenAI FunctionDeclaration export declaration = read_note_pure.to_genai_declaration() assert declaration.name == "read_note_pure" diff --git a/tests/unit/tools/test_read_note_tool.py b/tests/unit/tools/test_read_note_tool.py index f1e3d64..9b96f72 100644 --- a/tests/unit/tools/test_read_note_tool.py +++ b/tests/unit/tools/test_read_note_tool.py @@ -16,7 +16,7 @@ from unittest.mock import MagicMock, patch from artemis.context import ArtemisContext from artemis.graph.state import State -from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.base import ArtemisTool from artemis.tools.scratchpad import ( ReadNote, ReadNoteArgs, @@ -39,8 +39,8 @@ def mock_ctx(tmp_path): return ctx -def test_read_note_tool_subclass_and_registry(): - """Verify ReadNoteTool is a subclass of ArtemisTool and properly registered.""" +def test_read_note_tool_subclass(): + """Verify ReadNoteTool is a subclass of ArtemisTool.""" assert issubclass(ReadNoteTool, ArtemisTool) assert issubclass(ReadNote, ArtemisTool) assert isinstance(read_note, ArtemisTool) @@ -50,11 +50,6 @@ def test_read_note_tool_subclass_and_registry(): assert read_note.category == "memory" assert read_note.args_schema == ReadNoteArgs - # Registry lookup - reg_tool = ToolRegistry.get("read_note") - assert reg_tool is not None - assert isinstance(reg_tool, ReadNoteTool) - # GenAI FunctionDeclaration export declaration = read_note.to_genai_declaration() assert declaration.name == "read_note" diff --git a/tests/unit/tools/test_save_note_pure_tool.py b/tests/unit/tools/test_save_note_pure_tool.py index 1a91c62..982e719 100644 --- a/tests/unit/tools/test_save_note_pure_tool.py +++ b/tests/unit/tools/test_save_note_pure_tool.py @@ -16,7 +16,7 @@ from pathlib import Path from unittest.mock import MagicMock, patch from artemis.context import ArtemisContext -from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.base import ArtemisTool from artemis.tools.scratchpad import ( SaveNoteArgs, SaveNotePure, @@ -35,8 +35,8 @@ def mock_ctx(tmp_path): return ctx -def test_save_note_pure_tool_subclass_and_registry(): - """Verify SaveNotePureTool is a subclass of ArtemisTool and properly registered.""" +def test_save_note_pure_tool_subclass(): + """Verify SaveNotePureTool is a subclass of ArtemisTool.""" assert issubclass(SaveNotePureTool, ArtemisTool) assert issubclass(SaveNotePure, ArtemisTool) assert isinstance(save_note_pure, ArtemisTool) @@ -46,11 +46,6 @@ def test_save_note_pure_tool_subclass_and_registry(): assert save_note_pure.category == "memory" assert save_note_pure.args_schema == SaveNoteArgs - # Registry lookup - reg_tool = ToolRegistry.get("save_note_pure") - assert reg_tool is not None - assert isinstance(reg_tool, SaveNotePureTool) - # GenAI FunctionDeclaration export declaration = save_note_pure.to_genai_declaration() assert declaration.name == "save_note_pure" diff --git a/tests/unit/tools/test_save_note_tool.py b/tests/unit/tools/test_save_note_tool.py index 94a31b2..a5086ba 100644 --- a/tests/unit/tools/test_save_note_tool.py +++ b/tests/unit/tools/test_save_note_tool.py @@ -17,7 +17,7 @@ from unittest.mock import MagicMock, patch from artemis.context import ArtemisContext from artemis.graph.state import State -from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.base import ArtemisTool from artemis.tools.scratchpad import ( SaveNote, SaveNoteArgs, @@ -39,8 +39,8 @@ def mock_ctx(tmp_path): return ctx -def test_save_note_tool_subclass_and_registry(): - """Verify SaveNoteTool is a subclass of ArtemisTool and properly registered in ToolRegistry.""" +def test_save_note_tool_subclass(): + """Verify SaveNoteTool is a subclass of ArtemisTool.""" assert issubclass(SaveNoteTool, ArtemisTool) assert issubclass(SaveNote, ArtemisTool) assert isinstance(save_note, ArtemisTool) @@ -50,11 +50,6 @@ def test_save_note_tool_subclass_and_registry(): assert save_note.category == "memory" assert save_note.args_schema == SaveNoteArgs - # Registry lookup - reg_tool = ToolRegistry.get("save_note") - assert reg_tool is not None - assert isinstance(reg_tool, SaveNoteTool) - # GenAI FunctionDeclaration export declaration = save_note.to_genai_declaration() assert declaration.name == "save_note" diff --git a/tests/unit/tools/test_search_logs.py b/tests/unit/tools/test_search_logs.py index 0facb93..e1c1d86 100644 --- a/tests/unit/tools/test_search_logs.py +++ b/tests/unit/tools/test_search_logs.py @@ -15,7 +15,7 @@ from unittest.mock import MagicMock, patch from artemis.context import ArtemisContext -from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.base import ArtemisTool from artemis.tools.mobile.search_logs import ( SearchLogs, SearchLogsArgs, @@ -147,8 +147,8 @@ def mock_ctx(): return MagicMock(spec=ArtemisContext) -def test_search_logs_tool_subclass_and_registry(): - """Verify SearchLogsTool is a subclass of ArtemisTool and properly registered.""" +def test_search_logs_tool_subclass(): + """Verify SearchLogsTool is a subclass of ArtemisTool.""" assert issubclass(SearchLogsTool, ArtemisTool) assert issubclass(SearchLogs, ArtemisTool) assert issubclass(SearchLogsToolAlias, ArtemisTool) @@ -159,11 +159,6 @@ def test_search_logs_tool_subclass_and_registry(): assert search_logs.category == "diagnostic" assert search_logs.args_schema == SearchLogsArgs - # Registry lookup - reg_tool = ToolRegistry.get("search_logs") - assert reg_tool is not None - assert isinstance(reg_tool, SearchLogsTool) - # GenAI FunctionDeclaration export declaration = search_logs.to_genai_declaration() assert declaration.name == "search_logs" diff --git a/tests/unit/tools/test_update_note_pure_tool.py b/tests/unit/tools/test_update_note_pure_tool.py index 8d34bf5..ee34f45 100644 --- a/tests/unit/tools/test_update_note_pure_tool.py +++ b/tests/unit/tools/test_update_note_pure_tool.py @@ -16,7 +16,7 @@ from pathlib import Path from unittest.mock import MagicMock, patch from artemis.context import ArtemisContext -from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.base import ArtemisTool from artemis.tools.scratchpad import ( UpdateNoteArgs, UpdateNotePure, @@ -36,8 +36,8 @@ def mock_ctx(tmp_path): return ctx -def test_update_note_pure_tool_subclass_and_registry(): - """Verify UpdateNotePureTool is a subclass of ArtemisTool and properly registered.""" +def test_update_note_pure_tool_subclass(): + """Verify UpdateNotePureTool is a subclass of ArtemisTool.""" assert issubclass(UpdateNotePureTool, ArtemisTool) assert issubclass(UpdateNotePure, ArtemisTool) assert isinstance(update_note_pure, ArtemisTool) @@ -47,11 +47,6 @@ def test_update_note_pure_tool_subclass_and_registry(): assert update_note_pure.category == "memory" assert update_note_pure.args_schema == UpdateNoteArgs - # Registry lookup - reg_tool = ToolRegistry.get("update_note_pure") - assert reg_tool is not None - assert isinstance(reg_tool, UpdateNotePureTool) - # GenAI FunctionDeclaration export declaration = update_note_pure.to_genai_declaration() assert declaration.name == "update_note_pure" diff --git a/tests/unit/tools/test_update_note_tool.py b/tests/unit/tools/test_update_note_tool.py index 3acad96..a51f35c 100644 --- a/tests/unit/tools/test_update_note_tool.py +++ b/tests/unit/tools/test_update_note_tool.py @@ -17,7 +17,7 @@ from unittest.mock import MagicMock, patch from artemis.context import ArtemisContext from artemis.graph.state import State -from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.base import ArtemisTool from artemis.tools.scratchpad import ( UpdateNote, UpdateNoteArgs, @@ -40,8 +40,8 @@ def mock_ctx(tmp_path): return ctx -def test_update_note_tool_subclass_and_registry(): - """Verify UpdateNoteTool is a subclass of ArtemisTool and properly registered.""" +def test_update_note_tool_subclass(): + """Verify UpdateNoteTool is a subclass of ArtemisTool.""" assert issubclass(UpdateNoteTool, ArtemisTool) assert issubclass(UpdateNote, ArtemisTool) assert isinstance(update_note, ArtemisTool) @@ -51,11 +51,6 @@ def test_update_note_tool_subclass_and_registry(): assert update_note.category == "memory" assert update_note.args_schema == UpdateNoteArgs - # Registry lookup - reg_tool = ToolRegistry.get("update_note") - assert reg_tool is not None - assert isinstance(reg_tool, UpdateNoteTool) - # GenAI FunctionDeclaration export declaration = update_note.to_genai_declaration() assert declaration.name == "update_note" diff --git a/tests/unit/tools/test_video_analyzer_pure_tool.py b/tests/unit/tools/test_video_analyzer_pure_tool.py index cb1e349..597076d 100644 --- a/tests/unit/tools/test_video_analyzer_pure_tool.py +++ b/tests/unit/tools/test_video_analyzer_pure_tool.py @@ -15,7 +15,7 @@ from unittest.mock import AsyncMock, MagicMock, patch from artemis.context import ArtemisContext -from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.base import ArtemisTool from artemis.tools.video_tool import ( VideoAnalyzerArgs, VideoAnalyzerPure, @@ -32,8 +32,8 @@ def mock_ctx(): return ctx -def test_video_analyzer_pure_tool_subclass_and_registry(): - """Verify VideoAnalyzerPureTool is a subclass of ArtemisTool and properly registered.""" +def test_video_analyzer_pure_tool_subclass(): + """Verify VideoAnalyzerPureTool is a subclass of ArtemisTool.""" assert issubclass(VideoAnalyzerPureTool, ArtemisTool) assert issubclass(VideoAnalyzerPure, ArtemisTool) assert isinstance(video_analyzer_pure, ArtemisTool) @@ -43,11 +43,6 @@ def test_video_analyzer_pure_tool_subclass_and_registry(): assert video_analyzer_pure.category == "perception" assert video_analyzer_pure.args_schema == VideoAnalyzerArgs - # Registry lookup - reg_tool = ToolRegistry.get("video_analyzer_pure") - assert reg_tool is not None - assert isinstance(reg_tool, VideoAnalyzerPureTool) - # GenAI FunctionDeclaration export declaration = video_analyzer_pure.to_genai_declaration() assert declaration.name == "video_analyzer_pure" diff --git a/tests/unit/tools/test_video_analyzer_tool.py b/tests/unit/tools/test_video_analyzer_tool.py index ea891b4..121c702 100644 --- a/tests/unit/tools/test_video_analyzer_tool.py +++ b/tests/unit/tools/test_video_analyzer_tool.py @@ -16,7 +16,7 @@ from unittest.mock import AsyncMock, MagicMock, patch from artemis.context import ArtemisContext from artemis.graph.state import State -from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.base import ArtemisTool from artemis.tools.video_tool import ( DIAGNOSER_VIDEO_ANALYZER_DOCSTRING, OPERATOR_VIDEO_ANALYZER_DOCSTRING, @@ -41,8 +41,8 @@ def mock_state(): return MagicMock(spec=State) -def test_video_analyzer_tool_subclass_and_registry(): - """Verify VideoAnalyzerTool is a subclass of ArtemisTool and properly registered.""" +def test_video_analyzer_tool_subclass(): + """Verify VideoAnalyzerTool is a subclass of ArtemisTool.""" assert issubclass(VideoAnalyzerTool, ArtemisTool) assert isinstance(video_analyzer, ArtemisTool) assert isinstance(video_analyzer, VideoAnalyzerTool) @@ -51,11 +51,6 @@ def test_video_analyzer_tool_subclass_and_registry(): assert video_analyzer.category == "perception" assert video_analyzer.args_schema == VideoAnalyzerArgs - # Registry lookup - reg_tool = ToolRegistry.get("video_analyzer") - assert reg_tool is not None - assert isinstance(reg_tool, VideoAnalyzerTool) - # GenAI FunctionDeclaration export declaration = video_analyzer.to_genai_declaration() assert declaration.name == "video_analyzer" diff --git a/tests/unit/tools/test_wait_tool.py b/tests/unit/tools/test_wait_tool.py index 92a006e..bf2a33e 100644 --- a/tests/unit/tools/test_wait_tool.py +++ b/tests/unit/tools/test_wait_tool.py @@ -17,7 +17,7 @@ import time from unittest.mock import AsyncMock, MagicMock from artemis.context import ArtemisContext -from artemis.tools.base import ArtemisTool, ToolRegistry +from artemis.tools.base import ArtemisTool from artemis.tools.wait_tool import ( Wait, WaitArgs, @@ -35,8 +35,8 @@ def mock_ctx(): return ctx -def test_wait_tool_subclass_and_registry(): - """Verify WaitTool is a subclass of ArtemisTool and properly registered.""" +def test_wait_tool_subclass(): + """Verify WaitTool is a subclass of ArtemisTool.""" assert issubclass(WaitTool, ArtemisTool) assert issubclass(Wait, ArtemisTool) assert isinstance(wait, ArtemisTool) @@ -46,11 +46,6 @@ def test_wait_tool_subclass_and_registry(): assert wait.category == "system" assert wait.args_schema == WaitArgs - # Registry lookup - reg_tool = ToolRegistry.get("wait") - assert reg_tool is not None - assert isinstance(reg_tool, WaitTool) - # GenAI FunctionDeclaration export declaration = wait.to_genai_declaration() assert declaration.name == "wait" diff --git a/tests/unit/utils/test_optional_ocr.py b/tests/unit/utils/test_optional_ocr.py index ca11614..d8aa0dc 100644 --- a/tests/unit/utils/test_optional_ocr.py +++ b/tests/unit/utils/test_optional_ocr.py @@ -21,7 +21,6 @@ from artemis.agents.explorer.explorer import Explorer from artemis.config.settings import Settings from artemis.graph.perception import perception_node from artemis.graph.state import State -from artemis.tools.base import ToolRegistry from artemis.tools.index import get_tools_from_wrappers from artemis.tools.mobile.ocr import ( ocr_recognition, @@ -126,20 +125,8 @@ def test_ocr_tool_not_exposed_when_unconfigured(): with patch("artemis.tools.mobile.ocr.is_ocr_configured", return_value=False): assert not ocr_recognition.is_available() - # ToolRegistry.list_tools filters out unavailable tools - available_tools = ToolRegistry.list_tools(available_only=True) - assert "ocr_recognition" not in [t.name for t in available_tools] - - # ToolRegistry.get_langchain_tools filters out unavailable tools - mock_ctx = MagicMock() - lc_tools = ToolRegistry.get_langchain_tools(mock_ctx, available_only=True) - assert "ocr_recognition" not in [t.name for t in lc_tools] - - # ToolRegistry.get_genai_declarations filters out unavailable tools - genai_tools = ToolRegistry.get_genai_declarations(available_only=True) - assert "ocr_recognition" not in [t.name for t in genai_tools] - # get_tools_from_wrappers excludes ocr_recognition_wrapper + mock_ctx = MagicMock() wrapped_tools = get_tools_from_wrappers(mock_ctx, [ocr_recognition_wrapper]) assert len(wrapped_tools) == 0