inspect-ai 0.3.58__py3-none-any.whl → 0.3.59__py3-none-any.whl

inspect_ai/solver/_task_state.py

@@ -2,8 +2,9 @@ from collections.abc import Sequence
 from contextvars import ContextVar
 from copy import deepcopy
 from dataclasses import dataclass
+from itertools import tee
 from random import Random
-from typing import Any, Type, Union, cast, overload
+from typing import Any, Iterable, SupportsIndex, Type, Union, cast, overload
 
 from pydantic_core import to_jsonable_python
 
@@ -15,9 +16,13 @@ from inspect_ai.model import (
     ModelOutput,
 )
 from inspect_ai.model._call_tools import tools_info
+from inspect_ai.model._chat_message import ChatMessageBase
 from inspect_ai.model._model import sample_total_tokens
+from inspect_ai.scorer._metric import Score
+from inspect_ai.scorer._target import Target
 from inspect_ai.tool import Tool, ToolChoice
 from inspect_ai.tool._tool_def import ToolDef
+from inspect_ai.util._limit import SampleLimitExceededError
 from inspect_ai.util._store import Store, store_jsonable
 from inspect_ai.util._store_model import SMT
 
@@ -136,6 +141,7 @@ class TaskState:
         epoch: int,
         input: str | list[ChatMessage],
         messages: list[ChatMessage],
+        target: Target = Target(""),
         choices: list[str] | None = [],
         output: ModelOutput | None = None,
         message_limit: int | None = None,
@@ -161,10 +167,13 @@ class TaskState:
         or `input_text` only
         """
 
+        self.target = target
+        """The scoring target for this `Sample`."""
+
         self.metadata = metadata
         """Metadata from the `Sample` for this `TaskState`"""
 
-        self.messages = messages
+        self._messages: list[ChatMessage] = ChatMessageList(messages)
         """
         Chat conversation history for sample.
 
@@ -189,9 +198,7 @@ class TaskState:
         """
 
         self._message_limit = message_limit
-        self._message_limit_exceeded = False
         self._token_limit = token_limit
-        self._token_limit_exceeded = False
         self._completed = completed
 
         """Store for shared data"""
@@ -202,6 +209,9 @@ class TaskState:
         else:
             self.choices = Choices([])
 
+        self.scores: dict[str, Score] | None = None
+        """Scores yielded by running task."""
+
     @property
     def model(self) -> ModelName:
         """Name of model being evaluated."""
@@ -254,6 +264,16 @@ class TaskState:
         else:
             raise ValueError("user_prompt requested from TaskState but none available")
 
+    @property
+    def messages(self) -> list[ChatMessage]:
+        """Messages in chat history"""
+        return self._messages
+
+    @messages.setter
+    def messages(self, messages: list[ChatMessage]) -> None:
+        """Set messages in chat history."""
+        self._messages = ChatMessageList(messages)
+
     @property
     def max_messages(self) -> int | None:
         """Deprecated (use message_limit)."""
@@ -300,40 +320,7 @@ class TaskState:
     @property
     def completed(self) -> bool:
         """Is the task completed."""
-        # update messages
-        from inspect_ai.log._samples import set_active_sample_total_messages
-        from inspect_ai.log._transcript import SampleLimitEvent, transcript
-
-        set_active_sample_total_messages(len(self.messages))
-
-        if self._completed:
-            return True
-        elif self.message_limit and len(self.messages) >= self.message_limit:
-            # log if this is the first time we hit this
-            if not self._message_limit_exceeded:
-                self._message_limit_exceeded = True
-                transcript()._event(
-                    SampleLimitEvent(
-                        type="message",
-                        message=f"Sample completed: exceeded message limit ({self.message_limit})",
-                        limit=self.message_limit,
-                    )
-                )
-            return True
-        elif self.token_limit and self.token_usage >= self.token_limit:
-            # log if this is the first time we hit this
-            if not self._token_limit_exceeded:
-                self._token_limit_exceeded = True
-                transcript()._event(
-                    SampleLimitEvent(
-                        type="token",
-                        message=f"Sample completed: exceeded token limit ({self.token_limit:,})",
-                        limit=self.token_limit,
-                    )
-                )
-            return True
-        else:
-            return False
+        return self._completed
 
     @completed.setter
     def completed(self, completed: bool) -> None:
@@ -413,3 +400,58 @@ def state_jsonable(state: TaskState | None = None) -> dict[str, Any]:
 def sample_jsonable(sample: Sample) -> dict[str, Any]:
     jsonable = to_jsonable_python(sample, exclude_none=True, fallback=lambda _x: None)
     return cast(dict[str, Any], deepcopy(jsonable))
+
+
+class ChatMessageList(list[ChatMessage]):
+    def __init__(self, iterable: Iterable[ChatMessage]):
+        items, length = self._iterable_length(iterable)
+        self._check_size(length)
+        super().__init__(items)
+
+    def _check_size(self, additional_items: int = 1) -> None:
+        from inspect_ai.log._samples import active_sample_message_limit
+
+        messages_limit = active_sample_message_limit()
+        if messages_limit is not None:
+            messages = len(self) + additional_items
+            if messages > messages_limit:
+                raise SampleLimitExceededError(
+                    "message", value=messages, limit=messages_limit
+                )
+
+    def append(self, item: ChatMessage) -> None:
+        self._check_size()
+        super().append(item)
+
+    def extend(self, items: Iterable[ChatMessage]) -> None:
+        items, length = self._iterable_length(items)
+        self._check_size(length)
+        super().extend(items)
+
+    def insert(self, index: SupportsIndex, item: ChatMessage) -> None:
+        self._check_size()
+        super().insert(index, item)
+
+    @overload
+    def __setitem__(self, index: SupportsIndex, item: ChatMessage) -> None: ...
+
+    @overload
+    def __setitem__(self, index: slice, item: Iterable[ChatMessage]) -> None: ...
+
+    def __setitem__(
+        self, index: SupportsIndex | slice, item: ChatMessage | Iterable[ChatMessage]
+    ) -> None:
+        if isinstance(index, slice) and not isinstance(item, ChatMessageBase):
+            item, length = self._iterable_length(item)
+            size_change = length - len(self[index])
+            if size_change > 0:
+                self._check_size(size_change)
+
+        super().__setitem__(index, item)  # type: ignore[assignment,index]
+
+    def _iterable_length(
+        self, items: Iterable[ChatMessage]
+    ) -> tuple[Iterable[ChatMessage], int]:
+        items, counter = tee(items)
+        length = sum(1 for _ in counter)
+        return items, length

inspect_ai/tool/__init__.py

@@ -12,6 +12,7 @@ from ._tool_call import (
     ToolCall,
     ToolCallContent,
     ToolCallError,
+    ToolCallModelInput,
     ToolCallView,
     ToolCallViewer,
 )
@@ -42,6 +43,7 @@ __all__ = [
     "ContentVideo",
     "ToolCall",
     "ToolCallContent",
+    "ToolCallModelInput",
     "ToolCallView",
     "ToolCallViewer",
     "ToolChoice",

inspect_ai/tool/_tool.py

@@ -24,7 +24,7 @@ from inspect_ai._util.registry import (
     registry_tag,
 )
 
-from ._tool_call import ToolCallViewer
+from ._tool_call import ToolCallModelInput, ToolCallViewer
 
 logger = getLogger(__name__)
 
@@ -112,6 +112,7 @@ def tool(
     *,
     name: str | None = None,
     viewer: ToolCallViewer | None = None,
+    model_input: ToolCallModelInput | None = None,
     parallel: bool = True,
     prompt: str | None = None,
 ) -> Callable[[Callable[P, Tool]], Callable[P, Tool]]: ...
@@ -122,6 +123,7 @@ def tool(
     *,
     name: str | None = None,
     viewer: ToolCallViewer | None = None,
+    model_input: ToolCallModelInput | None = None,
     parallel: bool = True,
     prompt: str | None = None,
 ) -> Callable[P, Tool] | Callable[[Callable[P, Tool]], Callable[P, Tool]]:
@@ -135,6 +137,8 @@ def tool(
             will be used as the name of the tool.
         viewer (ToolCallViewer | None): Provide a custom view
             of tool call and context.
+        model_input (ToolCallModelInput | None): Provide a custom
+            function for playing back tool results as model input.
         parallel (bool):
             Does this tool support parallel execution?
             (defaults to True).
@@ -176,6 +180,9 @@ def tool(
                         TOOL_PROMPT: prompt,
                         TOOL_PARALLEL: parallel,
                         TOOL_VIEWER: viewer,
+                        TOOL_MODEL_INPUT: (
+                            model_input or getattr(tool, TOOL_INIT_MODEL_INPUT, None)
+                        ),
                     },
                 ),
                 *args,
@@ -195,3 +202,7 @@ def tool(
 TOOL_PROMPT = "prompt"
 TOOL_PARALLEL = "parallel"
 TOOL_VIEWER = "viewer"
+TOOL_MODEL_INPUT = "model_input"
+
+
+TOOL_INIT_MODEL_INPUT = "__TOOL_INIT_MODEL_INPUT__"

inspect_ai/tool/_tool_call.py

@@ -3,6 +3,8 @@ from typing import Any, Callable, Literal
 
 from pydantic import BaseModel, Field
 
+from inspect_ai._util.content import Content
+
 
 class ToolCallContent(BaseModel):
     """Content to include in tool call view."""
@@ -71,3 +73,11 @@ class ToolCallError:
 
 ToolCallViewer = Callable[[ToolCall], ToolCallView]
 """Custom view renderer for tool calls."""
+
+
+ToolCallModelInput = Callable[[int, int, str | list[Content]], str | list[Content]]
+"""Determine how tool call results are played back as model input.
+
+The first argument is an index into the total number of tool results
+for this tool in the message history, the second is the total number.
+"""

inspect_ai/tool/_tool_def.py

@@ -13,8 +13,8 @@ from inspect_ai._util.registry import (
     set_registry_params,
 )
 
-from ._tool import TOOL_PARALLEL, TOOL_PROMPT, TOOL_VIEWER, Tool
-from ._tool_call import ToolCallViewer
+from ._tool import TOOL_MODEL_INPUT, TOOL_PARALLEL, TOOL_PROMPT, TOOL_VIEWER, Tool
+from ._tool_call import ToolCallModelInput, ToolCallViewer
 from ._tool_description import (
     ToolDescription,
     set_tool_description,
@@ -33,6 +33,7 @@ class ToolDef:
         parameters: dict[str, str] | ToolParams | None = None,
         parallel: bool | None = None,
         viewer: ToolCallViewer | None = None,
+        model_input: ToolCallModelInput | None = None,
     ) -> None:
         """Tool definition.
 
@@ -46,6 +47,8 @@ class ToolDef:
           parallel (bool | None): Does the tool support parallel execution
              (defaults to True if not specified)
           viewer (ToolCallViewer | None): Optional tool call viewer implementation.
+          model_input (ToolCallModelInput | None): Optional function that determines how
+              tool call results are played back as model input.
 
         Returns:
           Tool definition.
@@ -68,6 +71,7 @@ class ToolDef:
             parameters = parameters if parameters is not None else tdef.parameters
             self.parallel = parallel if parallel is not None else tdef.parallel
             self.viewer = viewer or tdef.viewer
+            self.model_input = model_input or tdef.model_input
 
         # if its not a tool then extract tool_info if all fields have not
         # been provided explicitly
@@ -97,6 +101,7 @@ class ToolDef:
             # behavioral attributes
             self.parallel = parallel is not False
             self.viewer = viewer
+            self.model_input = model_input
 
     tool: Callable[..., Any]
     """Callable to execute tool."""
@@ -116,6 +121,9 @@ class ToolDef:
     viewer: ToolCallViewer | None
     """Custom viewer for tool call"""
 
+    model_input: ToolCallModelInput | None
+    """Custom model input presenter for tool calls."""
+
     def as_tool(self) -> Tool:
         """Convert a ToolDef to a Tool."""
         tool = self.tool
@@ -159,11 +167,12 @@ class ToolDefFields(NamedTuple):
     parameters: ToolParams
     parallel: bool
     viewer: ToolCallViewer | None
+    model_input: ToolCallModelInput | None
 
 
 def tool_def_fields(tool: Tool) -> ToolDefFields:
     # get tool_info
-    name, prompt, parallel, viewer = tool_registry_info(tool)
+    name, prompt, parallel, viewer, model_input = tool_registry_info(tool)
     tool_info = parse_tool_info(tool)
 
     # if there is a description then append any prompt to the
@@ -213,15 +222,17 @@ def tool_def_fields(tool: Tool) -> ToolDefFields:
         parameters=tool_info.parameters,
         parallel=parallel,
         viewer=viewer,
+        model_input=model_input,
     )
 
 
 def tool_registry_info(
     tool: Tool,
-) -> tuple[str, str | None, bool, ToolCallViewer | None]:
+) -> tuple[str, str | None, bool, ToolCallViewer | None, ToolCallModelInput | None]:
     info = registry_info(tool)
     name = info.name.split("/")[-1]
     prompt = info.metadata.get(TOOL_PROMPT, None)
     parallel = info.metadata.get(TOOL_PARALLEL, True)
     viewer = info.metadata.get(TOOL_VIEWER, None)
-    return name, prompt, parallel, viewer
+    model_input = info.metadata.get(TOOL_MODEL_INPUT, None)
+    return name, prompt, parallel, viewer, model_input

inspect_ai/tool/_tool_with.py

@@ -1,4 +1,4 @@
-from copy import copy
+from copy import deepcopy
 
 from inspect_ai._util.registry import (
     registry_info,
@@ -6,8 +6,9 @@ from inspect_ai._util.registry import (
     set_registry_info,
     set_registry_params,
 )
+from inspect_ai.tool._tool_call import ToolCallModelInput, ToolCallViewer
 
-from ._tool import Tool
+from ._tool import TOOL_MODEL_INPUT, TOOL_PARALLEL, TOOL_VIEWER, Tool
 from ._tool_description import ToolDescription, set_tool_description
 from ._tool_info import parse_tool_info
 
@@ -17,6 +18,9 @@ def tool_with(
     name: str | None = None,
     description: str | None = None,
     parameters: dict[str, str] | None = None,
+    parallel: bool | None = None,
+    viewer: ToolCallViewer | None = None,
+    model_input: ToolCallModelInput | None = None,
 ) -> Tool:
     """Tool with modifications to name and descriptions.
 
@@ -25,6 +29,11 @@ def tool_with(
        name (str | None): Tool name (optional).
        description (str | None): Tool description (optional).
        parameters (dict[str,str] | None): Parameter descriptions (optional)
+       parallel (bool | None): Does the tool support parallel execution
+          (defaults to True if not specified)
+       viewer (ToolCallViewer | None): Optional tool call viewer implementation.
+       model_input (ToolCallModelInput | None): Optional function that determines how
+           tool call results are played back as model input.
 
     Returns:
        A copy of the passed tool with the specified descriptive information.
@@ -46,8 +55,16 @@ def tool_with(
             ]
 
     # copy the tool and set the descriptions on the new copy
-    tool_copy = copy(tool)
-    set_registry_info(tool_copy, registry_info(tool))
+    tool_copy = deepcopy(tool)
+    info = registry_info(tool).model_copy()
+    if parallel is not None:
+        info.metadata[TOOL_PARALLEL] = parallel
+    elif viewer is not None:
+        info.metadata[TOOL_VIEWER] = viewer
+    elif model_input is not None:
+        info.metadata[TOOL_MODEL_INPUT] = model_input
+
+    set_registry_info(tool_copy, info)
     set_registry_params(tool_copy, registry_params(tool))
     set_tool_description(
         tool_copy,

inspect_ai/tool/beta/__init__.py

@@ -0,0 +1,5 @@
+from ._computer import computer
+
+__all__ = [
+    "computer",
+]

inspect_ai/tool/beta/_computer/__init__.py

@@ -0,0 +1,3 @@
+from ._computer import computer
+
+__all__ = ["computer"]

inspect_ai/tool/beta/_computer/_common.py

@@ -0,0 +1,133 @@
+import json
+from textwrap import dedent
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+from inspect_ai._util.content import ContentText
+from inspect_ai._util.error import PrerequisiteError
+from inspect_ai.model import ContentImage
+from inspect_ai.tool import ToolError, ToolResult
+from inspect_ai.util._sandbox.context import sandbox_with
+from inspect_ai.util._sandbox.environment import SandboxEnvironment
+
+Action = Literal[
+    "key",
+    "type",
+    "mouse_move",
+    "left_click",
+    "left_click_drag",
+    "right_click",
+    "middle_click",
+    "double_click",
+    "screenshot",
+    "cursor_position",
+]
+
+
+class ToolExecResult(BaseModel):
+    output: str | None = Field(default=None)
+    error: str | None = Field(default=None)
+    base64_image: str | None = Field(default=None)
+
+
+async def _send_cmd(cmdTail: list[str], timeout: int | None = None) -> ToolResult:
+    from inspect_ai.log._samples import sample_active
+
+    sample = sample_active()
+    assert sample
+    sample_id = sample.sample.id
+    assert sample_id
+
+    cmd = ["python3", "/opt/inspect/tool/computer_tool.py", "--action"] + cmdTail
+
+    raw_exec_result = await (await computer_sandbox()).exec(cmd, timeout=timeout)
+
+    if not raw_exec_result.success:
+        raise RuntimeError(
+            f"Failure executing command: ${cmd} {raw_exec_result.stderr}"
+        )
+
+    result = ToolExecResult(**json.loads(raw_exec_result.stdout))
+
+    if result.error:
+        raise ToolError(result.error)
+
+    image = (
+        ContentImage(image=f"data:image/png;base64,{result.base64_image}")
+        if result.base64_image
+        else None
+    )
+    text = result.output if result.output and len(result.output) > 0 else None
+
+    if text is not None and image is not None:
+        return [ContentText(text=text), image]
+
+    if text is not None:
+        return text
+
+    if image is not None:
+        return [image]
+
+    return "OK"
+
+
+async def cursor_position(timeout: int | None = None) -> ToolResult:
+    return await _send_cmd(["cursor_position"], timeout=timeout)
+
+
+async def screenshot(timeout: int | None = None) -> ToolResult:
+    return await _send_cmd(["screenshot"], timeout=timeout)
+
+
+async def mouse_move(x: int, y: int, timeout: int | None = None) -> ToolResult:
+    return await _send_cmd(
+        ["mouse_move", "--coordinate", f"{x}", f"{y}"], timeout=timeout
+    )
+
+
+async def left_click(timeout: int | None = None) -> ToolResult:
+    return await _send_cmd(["left_click"], timeout=timeout)
+
+
+async def left_click_drag(x: int, y: int, timeout: int | None = None) -> ToolResult:
+    return await _send_cmd(
+        ["left_click_drag", "--coordinate", f"{x}", f"{y}"], timeout=timeout
+    )
+
+
+async def right_click(timeout: int | None = None) -> ToolResult:
+    return await _send_cmd(["right_click"], timeout=timeout)
+
+
+async def middle_click(timeout: int | None = None) -> ToolResult:
+    return await _send_cmd(["middle_click"], timeout=timeout)
+
+
+async def double_click(timeout: int | None = None) -> ToolResult:
+    return await _send_cmd(["double_click"], timeout=timeout)
+
+
+async def press_key(key: str, timeout: int | None = None) -> ToolResult:
+    return await _send_cmd(["key", "--text", key], timeout=timeout)
+
+
+async def type(text: str, timeout: int | None = None) -> ToolResult:
+    return await _send_cmd(["type", "--text", text], timeout=timeout)
+
+
+async def computer_sandbox() -> SandboxEnvironment:
+    sb = await sandbox_with("/opt/inspect/tool/computer_tool.py")
+    if sb:
+        return sb
+    else:
+        raise PrerequisiteError(
+            dedent("""
+                The computer tool service was not found in any of the sandboxes for this sample. Please add the computer tool service to your configuration. For example, the following Docker compose file uses the aisiuk/inspect-computer-tool:latest-beta image as its default sandbox:
+
+                services:
+                  default:
+                    image: "aisiuk/inspect-computer-tool:latest-beta"
+                    init: true
+                """).strip()
+        )