398 lines
14 KiB
Python
398 lines
14 KiB
Python
# 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.
|
|
|
|
"""Shared history search, step replay, and screenshot tools.
|
|
|
|
Each tool's args_schema defines both its LangChain interface for Pro and its
|
|
ToolDeclaration for Flash. The MCP trace inspector calls the underlying
|
|
functions with an OfflineHistoryReader.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Literal
|
|
|
|
from langchain_core.tools import BaseTool
|
|
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
|
|
from artemis.tools.history.replay import (
|
|
DEFAULT_MAX_STEPS as DEFAULT_REPLAY_MAX_STEPS,
|
|
)
|
|
from artemis.tools.history.replay import (
|
|
DEFAULT_MAX_TOKENS as DEFAULT_REPLAY_MAX_TOKENS,
|
|
)
|
|
from artemis.tools.history.replay import replay_steps_text
|
|
from artemis.tools.history.screenshot import ScreenshotResult, load_step_screenshot
|
|
from artemis.tools.history.search import search_history_text
|
|
from artemis.tools.tool_wrapper import ToolWrapper
|
|
from artemis.utils.logger import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
__all__ = [
|
|
"HISTORY_TOOLS",
|
|
"HISTORY_TOOL_DECLARATIONS",
|
|
"HISTORY_TOOL_NAMES",
|
|
"HISTORY_TOOL_WRAPPERS",
|
|
"GET_STEP_SCREENSHOT_TOOL",
|
|
"REPLAY_STEPS_TOOL",
|
|
"SEARCH_HISTORY_TOOL",
|
|
"GetStepScreenshotArgs",
|
|
"GetStepScreenshotTool",
|
|
"ReplayStepsArgs",
|
|
"ReplayStepsTool",
|
|
"ScreenshotResult",
|
|
"SearchHistoryArgs",
|
|
"SearchHistoryTool",
|
|
"get_history_tools",
|
|
"get_replay_steps_tool",
|
|
"get_search_history_tool",
|
|
"get_step_screenshot",
|
|
"get_step_screenshot_tool",
|
|
"history_available",
|
|
"history_tool_by_name",
|
|
"history_tool_declarations",
|
|
"load_step_screenshot",
|
|
"replay_steps",
|
|
"replay_steps_text",
|
|
"search_history",
|
|
"search_history_available",
|
|
"search_history_text",
|
|
]
|
|
|
|
|
|
# --- Configuration -----------------------------------------------------------------------
|
|
|
|
|
|
def _memory_config(warn: bool = True) -> Any:
|
|
try:
|
|
from artemis.config import load_agent_config
|
|
|
|
return load_agent_config().memory
|
|
except Exception as e:
|
|
if warn:
|
|
logger.debug(f"memory config unavailable, using defaults: {e}")
|
|
return None
|
|
|
|
|
|
def _recall_config(warn: bool = True) -> Any:
|
|
return getattr(_memory_config(warn), "recall", None)
|
|
|
|
|
|
def _replay_config(warn: bool = True) -> Any:
|
|
return getattr(_memory_config(warn), "replay", None)
|
|
|
|
|
|
def history_available(ctx: Any) -> bool:
|
|
"""History tools need a stored history to read (a DataEngine session)."""
|
|
return ctx is not None and getattr(ctx, "data_engine", None) is not None
|
|
|
|
|
|
def search_history_available(ctx: Any) -> bool:
|
|
"""``search_history`` is additionally gated by ``agent.memory.recall.enabled``."""
|
|
if not history_available(ctx):
|
|
return False
|
|
return bool(getattr(_recall_config(warn=False), "enabled", True))
|
|
|
|
|
|
# --- search_history ----------------------------------------------------------------------
|
|
|
|
|
|
class SearchHistoryArgs(BaseModel):
|
|
"""Arguments schema for the search_history tool."""
|
|
|
|
query: str = Field(
|
|
default="",
|
|
description=(
|
|
"Keywords to search for across the whole execution history"
|
|
" (screen descriptions, actions, results, reasoning, tool calls,"
|
|
" notes, on-screen text, package/activity names). Case-insensitive;"
|
|
" every whitespace-separated term is matched independently. May be"
|
|
" empty when step_range is given."
|
|
),
|
|
)
|
|
step_range: list[int] | None = Field(
|
|
default=None,
|
|
description=(
|
|
"Optional [start, end] step-number range (inclusive) to restrict the"
|
|
" search — e.g. the range shown in a compressed-history block. Also"
|
|
" returns the full per-step action ledger of the range."
|
|
),
|
|
)
|
|
max_results: int = Field(
|
|
default=5,
|
|
ge=1,
|
|
description="Maximum results to return (server-side cap still applies).",
|
|
)
|
|
|
|
|
|
SEARCH_HISTORY_DESCRIPTION = (
|
|
"Deterministic keyword / step-range lookup over the full stored execution"
|
|
" history: step summaries, exact actions and results, reasoning, tool"
|
|
" calls, notes, on-screen text (OCR/UI tree, incl. package/activity names)"
|
|
" and compressed-history ledgers. Every hit carries its step number; use"
|
|
" `replay_steps` for the full record of a step and `get_step_screenshot`"
|
|
" for its image."
|
|
)
|
|
|
|
|
|
class SearchHistoryTool(ArtemisTool):
|
|
"""Deterministic search over the stored execution history."""
|
|
|
|
def __init__(self):
|
|
super().__init__(
|
|
name="search_history",
|
|
description=SEARCH_HISTORY_DESCRIPTION,
|
|
args_schema=SearchHistoryArgs,
|
|
category="memory",
|
|
)
|
|
|
|
def is_available(self, ctx: Any = None) -> bool:
|
|
return search_history_available(ctx)
|
|
|
|
# pylint: disable=too-many-arguments
|
|
async def execute(
|
|
self,
|
|
driver: Any = None, # pylint: disable=unused-argument
|
|
ctx: Any = None,
|
|
query: str = "",
|
|
step_range: list[int] | None = None,
|
|
max_results: int = 5,
|
|
**kwargs: Any,
|
|
) -> str:
|
|
reader = getattr(ctx, "data_engine", None) if ctx else None
|
|
if reader is None:
|
|
return "search_history unavailable: no active execution history."
|
|
if not query and not step_range:
|
|
return (
|
|
"search_history needs a query and/or a step_range — e.g."
|
|
' search_history(query="login timeout") or'
|
|
' search_history(query="", step_range=[1, 40]).'
|
|
)
|
|
try:
|
|
return search_history_text(
|
|
reader,
|
|
query=query or "",
|
|
step_range=step_range,
|
|
max_results=int(max_results or 5),
|
|
recall_config=_recall_config(),
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"search_history failed: {e}")
|
|
return f"search_history failed: {e}"
|
|
|
|
|
|
# --- replay_steps ------------------------------------------------------------------------
|
|
|
|
|
|
class ReplayStepsArgs(BaseModel):
|
|
"""Arguments schema for the replay_steps tool."""
|
|
|
|
start_step: int = Field(..., description="First step to replay (1-based step number).")
|
|
end_step: int | None = Field(
|
|
default=None,
|
|
description=(
|
|
"Last step of the inclusive range to replay; omit to replay the single"
|
|
" start_step. Reversed bounds are swapped."
|
|
),
|
|
)
|
|
|
|
|
|
REPLAY_STEPS_DESCRIPTION = (
|
|
"Replays recorded steps exactly as the executing agent's own context showed"
|
|
" them: what the screen showed, its reasoning, every tool call (name,"
|
|
" arguments, result), the planned action, any pre-execution interception"
|
|
" and the execution result. Coordinates normalized [x, y]. Up to"
|
|
f" {DEFAULT_REPLAY_MAX_STEPS} steps per call. Screenshots are not included"
|
|
" — fetch them with `get_step_screenshot`."
|
|
)
|
|
|
|
|
|
class ReplayStepsTool(ArtemisTool):
|
|
"""Full replay of a range of recorded steps."""
|
|
|
|
def __init__(self):
|
|
super().__init__(
|
|
name="replay_steps",
|
|
description=REPLAY_STEPS_DESCRIPTION,
|
|
args_schema=ReplayStepsArgs,
|
|
category="memory",
|
|
)
|
|
|
|
def is_available(self, ctx: Any = None) -> bool:
|
|
return history_available(ctx)
|
|
|
|
async def execute(
|
|
self,
|
|
driver: Any = None, # pylint: disable=unused-argument
|
|
ctx: Any = None,
|
|
start_step: Any = None,
|
|
end_step: Any = None,
|
|
**kwargs: Any,
|
|
) -> str:
|
|
reader = getattr(ctx, "data_engine", None) if ctx else None
|
|
if reader is None:
|
|
return "Error: no execution history available."
|
|
cfg = _replay_config()
|
|
try:
|
|
return replay_steps_text(
|
|
reader,
|
|
start_step,
|
|
end_step,
|
|
max_steps=int(getattr(cfg, "max_steps", None) or DEFAULT_REPLAY_MAX_STEPS),
|
|
max_tokens=int(getattr(cfg, "max_tokens", None) or DEFAULT_REPLAY_MAX_TOKENS),
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"replay_steps failed: {e}")
|
|
return f"replay_steps failed: {e}"
|
|
|
|
|
|
# --- get_step_screenshot -----------------------------------------------------------------
|
|
|
|
|
|
class GetStepScreenshotArgs(BaseModel):
|
|
"""Arguments schema for the get_step_screenshot tool."""
|
|
|
|
step_number: int = Field(..., description="The step whose screenshot to attach (1-based).")
|
|
which: Literal["pre", "post", "overlay"] = Field(
|
|
default="pre",
|
|
description=(
|
|
"'pre' = screen observed at the start of the step; 'post' = screen"
|
|
" after the action (when recorded); 'overlay' = the pre screenshot"
|
|
" with the step's action drawn on it."
|
|
),
|
|
)
|
|
|
|
|
|
GET_STEP_SCREENSHOT_DESCRIPTION = (
|
|
"Attaches one recorded screenshot of a step: 'pre' (screen observed at the"
|
|
" start of the step), 'post' (screen after the action, when recorded) or"
|
|
" 'overlay' (the pre screenshot with the step's action drawn on it — tap"
|
|
" point / swipe path — to check the action landed on the intended element)."
|
|
)
|
|
|
|
|
|
class GetStepScreenshotTool(ArtemisTool):
|
|
"""One recorded step screenshot, attached to the conversation."""
|
|
|
|
def __init__(self):
|
|
super().__init__(
|
|
name="get_step_screenshot",
|
|
description=GET_STEP_SCREENSHOT_DESCRIPTION,
|
|
args_schema=GetStepScreenshotArgs,
|
|
category="memory",
|
|
)
|
|
|
|
def is_available(self, ctx: Any = None) -> bool:
|
|
return history_available(ctx)
|
|
|
|
async def execute(
|
|
self,
|
|
driver: Any = None, # pylint: disable=unused-argument
|
|
ctx: Any = None,
|
|
step_number: Any = None,
|
|
which: str = "pre",
|
|
**kwargs: Any,
|
|
) -> list[dict[str, Any]] | str:
|
|
reader = getattr(ctx, "data_engine", None) if ctx else None
|
|
if reader is None:
|
|
return "Error: no execution history available."
|
|
try:
|
|
return load_step_screenshot(reader, step_number, which).to_content_blocks()
|
|
except Exception as e:
|
|
logger.error(f"get_step_screenshot failed: {e}")
|
|
return f"get_step_screenshot failed: {e}"
|
|
|
|
|
|
# --- Instances, declarations, exports ------------------------------------------------------
|
|
|
|
search_history = SearchHistoryTool()
|
|
replay_steps = ReplayStepsTool()
|
|
get_step_screenshot = GetStepScreenshotTool()
|
|
|
|
HISTORY_TOOLS: tuple[ArtemisTool, ...] = (search_history, replay_steps, get_step_screenshot)
|
|
HISTORY_TOOL_NAMES: frozenset[str] = frozenset(t.name for t in HISTORY_TOOLS)
|
|
|
|
#: Flash-profile declarations, derived from the very same args schemas.
|
|
SEARCH_HISTORY_TOOL: ToolDeclaration = search_history.to_tool_declaration()
|
|
REPLAY_STEPS_TOOL: ToolDeclaration = replay_steps.to_tool_declaration()
|
|
GET_STEP_SCREENSHOT_TOOL: ToolDeclaration = get_step_screenshot.to_tool_declaration()
|
|
HISTORY_TOOL_DECLARATIONS: tuple[ToolDeclaration, ...] = (
|
|
SEARCH_HISTORY_TOOL,
|
|
REPLAY_STEPS_TOOL,
|
|
GET_STEP_SCREENSHOT_TOOL,
|
|
)
|
|
|
|
|
|
def history_tool_by_name(name: str) -> ArtemisTool | None:
|
|
"""The shared tool instance behind a history tool name (Flash dispatch)."""
|
|
raw = name.split(":")[-1] if ":" in name else name
|
|
return next((t for t in HISTORY_TOOLS if t.name == raw), None)
|
|
|
|
|
|
def history_tool_declarations(ctx: Any) -> list[ToolDeclaration]:
|
|
"""Declarations of the history tools available for ``ctx`` (Flash runner)."""
|
|
return [tool.to_tool_declaration() for tool in HISTORY_TOOLS if tool.is_available(ctx)]
|
|
|
|
|
|
def get_search_history_tool(ctx: Any) -> BaseTool:
|
|
"""Exports search_history as a traced LangChain BaseTool."""
|
|
return trace_langchain_tool(search_history.to_langchain_tool(ctx), ctx)
|
|
|
|
|
|
def get_replay_steps_tool(ctx: Any) -> BaseTool:
|
|
"""Exports replay_steps as a traced LangChain BaseTool."""
|
|
return trace_langchain_tool(replay_steps.to_langchain_tool(ctx), ctx)
|
|
|
|
|
|
def get_step_screenshot_tool(ctx: Any) -> BaseTool:
|
|
"""Exports get_step_screenshot as a traced LangChain BaseTool."""
|
|
return trace_langchain_tool(get_step_screenshot.to_langchain_tool(ctx), ctx)
|
|
|
|
|
|
def get_history_tools(ctx: Any) -> list[BaseTool]:
|
|
"""All three history tools as LangChain tools (read-only agents mount them as a set)."""
|
|
return [get_search_history_tool(ctx), get_replay_steps_tool(ctx), get_step_screenshot_tool(ctx)]
|
|
|
|
|
|
search_history_wrapper = ToolWrapper(
|
|
tool_fn_getter=get_search_history_tool,
|
|
on_success_fn=lambda *a, **k: "Searched history",
|
|
on_failure_fn=lambda err: f"search_history failed: {err}",
|
|
is_available_fn=search_history_available,
|
|
)
|
|
|
|
replay_steps_wrapper = ToolWrapper(
|
|
tool_fn_getter=get_replay_steps_tool,
|
|
on_success_fn=lambda *a, **k: "Replayed steps",
|
|
on_failure_fn=lambda err: f"replay_steps failed: {err}",
|
|
is_available_fn=history_available,
|
|
)
|
|
|
|
get_step_screenshot_wrapper = ToolWrapper(
|
|
tool_fn_getter=get_step_screenshot_tool,
|
|
on_success_fn=lambda *a, **k: "Fetched step screenshot",
|
|
on_failure_fn=lambda err: f"get_step_screenshot failed: {err}",
|
|
is_available_fn=history_available,
|
|
)
|
|
|
|
HISTORY_TOOL_WRAPPERS: tuple[ToolWrapper, ...] = (
|
|
search_history_wrapper,
|
|
replay_steps_wrapper,
|
|
get_step_screenshot_wrapper,
|
|
)
|