inspect-ai 0.3.91__py3-none-any.whl → 0.3.93__py3-none-any.whl

inspect_ai/solver/_basic_agent.py

@@ -13,8 +13,8 @@ from inspect_ai.scorer._score import score
 from inspect_ai.solver._chain import chain
 from inspect_ai.tool._tool import Tool, ToolResult, tool
 from inspect_ai.tool._tool_with import tool_with
+from inspect_ai.util._limit import token_limit as create_token_limit
 
-from ._limit import SampleLimitExceededError
 from ._prompt import system_message
 from ._solver import Generate, Solver, solver
 from ._task_state import TaskState
@@ -172,14 +172,11 @@ def basic_agent(
             # (if there is no message_limit then default to 50)
             state.message_limit = message_limit or state.message_limit or 50
 
-            # resolve token limit
-            state.token_limit = token_limit or state.token_limit
-
             # track attempts
             attempts = 0
 
-            try:
-                # main loop (state.completed checks message_limit and token_limit)
+            with create_token_limit(token_limit):
+                # main loop
                 while not state.completed:
                     # generate output and append assistant message
                     state.output = await get_model().generate(
@@ -247,10 +244,6 @@ def basic_agent(
                     else:
                         state.messages.append(ChatMessageUser(content=continue_message))
 
-            # propagate current state along with sample limit exceeded
-            except SampleLimitExceededError as ex:
-                raise ex.with_state(state)
-
             return state
 
         return solve

inspect_ai/solver/_task_state.py

@@ -2,9 +2,8 @@ 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, Iterable, SupportsIndex, Type, Union, cast, overload
+from typing import Any, Type, Union, cast, overload
 
 from pydantic_core import to_jsonable_python
 from shortuuid import uuid
@@ -18,12 +17,18 @@ 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 (
+    check_message_limit,
+    check_token_limit,
+)
+from inspect_ai.util._limit import message_limit as create_message_limit
+from inspect_ai.util._limit import token_limit as create_token_limit
+from inspect_ai.util._limited_conversation import ChatMessageList
 from inspect_ai.util._store import Store, store_jsonable
 from inspect_ai.util._store_model import SMT
 
@@ -159,11 +164,11 @@ class TaskState:
         self._input = input
         self._target = target
         self._metadata = metadata
-        self._messages: list[ChatMessage] = ChatMessageList(messages, self)
+        self._messages: list[ChatMessage] = ChatMessageList(messages)
         self._tools: list[Tool] = []
         self._output = output if output else ModelOutput(model=str(model))
-        self._message_limit = message_limit
-        self._token_limit = token_limit
+        self._message_limit = create_message_limit(message_limit)
+        self._token_limit = create_token_limit(token_limit)
         self._completed = completed
         self._store = Store()
         self._uuid = uuid()
@@ -254,7 +259,7 @@ class TaskState:
 
     @messages.setter
     def messages(self, messages: list[ChatMessage]) -> None:
-        self._messages = ChatMessageList(messages, self)
+        self._messages = ChatMessageList(messages)
 
     @property
     def output(self) -> ModelOutput:
@@ -302,12 +307,16 @@ class TaskState:
     @property
     def message_limit(self) -> int | None:
         """Limit on total messages allowed per conversation."""
-        return self._message_limit
+        return self._message_limit.limit
 
     @message_limit.setter
     def message_limit(self, messages: int | None) -> None:
-        """Set limit on total messages allowed per conversation."""
-        self._message_limit = messages
+        """Set limit on total messages allowed per conversation.
+
+        Also checks whether the current message count exceeds the new limit.
+        """
+        self._message_limit.limit = messages
+        check_message_limit(len(self.messages), raise_for_equal=False)
 
         from inspect_ai.log._samples import set_active_sample_message_limit
 
@@ -316,12 +325,16 @@ class TaskState:
     @property
     def token_limit(self) -> int | None:
         """Limit on total tokens allowed per conversation."""
-        return self._token_limit
+        return self._token_limit.limit
 
     @token_limit.setter
     def token_limit(self, tokens: int | None) -> None:
-        """Set limit on total tokens allowed per conversation."""
-        self._token_limit = tokens
+        """Set limit on total tokens allowed per conversation.
+
+        Also checks whether the current token usage exceeds the new limit.
+        """
+        self._token_limit.limit = tokens
+        check_token_limit()
 
         from inspect_ai.log._samples import set_active_sample_token_limit
 
@@ -340,24 +353,11 @@ class TaskState:
         """
         from inspect_ai.log._samples import set_active_sample_total_messages
 
-        from ._limit import SampleLimitExceededError
-
         # update messages
         set_active_sample_total_messages(len(self.messages))
 
         if self._completed:
             return True
-        elif self.message_limit and len(self.messages) >= self.message_limit:
-            raise SampleLimitExceededError(
-                "message",
-                value=len(self.messages),
-                limit=self.message_limit,
-                state=self,
-            )
-        elif self.token_limit and self.token_usage >= self.token_limit:
-            raise SampleLimitExceededError(
-                "token", value=self.token_usage, limit=self.token_limit, state=self
-            )
         else:
             check_sample_interrupt()
             return self._completed
@@ -445,65 +445,3 @@ 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], parent_state: TaskState):
-        self.parent_state = parent_state
-        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
-
-        from ._limit import SampleLimitExceededError
-
-        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,
-                    message=None,
-                    state=self.parent_state,
-                )
-
-    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/_json_rpc_helpers.py

@@ -4,7 +4,7 @@ from typing import Literal, Protocol, Type, TypeAlias, TypeVar
 
 from pydantic import BaseModel, RootModel
 
-from inspect_ai.tool._tool import ToolError
+from inspect_ai.tool._tool import ToolError, ToolParsingError
 
 
 class JSONRPCResponseBase(BaseModel):
@@ -70,6 +70,7 @@ async def exec_scalar_request(
     params: JSONRPCParamsType,
     result_type: Type[ScalarT],
     transport: JSONRPCTransport,
+    server_error_mapper: JSONRPCServerErrorMapper,
 ) -> ScalarT:
     """
     Execute a JSON-RPC command expecting a scalar result.
@@ -79,6 +80,7 @@ async def exec_scalar_request(
       params (JSONRPCParamsType): The parameters for the JSON-RPC method.
       result_type (Type[ScalarT]): The scalar type (str, int, float, bool, None) to validate the result against.
       transport (JSONRPCTransport): The transport callable to use for the RPC communication.
+      server_error_mapper (JSONRPCServerErrorMapper): A callable to map server specific JSON-RPC errors to exceptions.
 
     Returns:
       ScalarT: The scalar result of the JSON-RPC call.
@@ -88,7 +90,12 @@ async def exec_scalar_request(
       ToolParsingError: If the JSON-RPC response contains a specific error code indicating a parsing error.
       ValueError: If the result is not of the expected scalar type.
     """
-    rpc_result = await _exec_request(method=method, params=params, transport=transport)
+    rpc_result = await _exec_request(
+        method=method,
+        params=params,
+        transport=transport,
+        server_error_mapper=server_error_mapper,
+    )
     if (result_type is type(None) and rpc_result is not None) or not isinstance(
         rpc_result, result_type
     ):
@@ -101,6 +108,7 @@ async def exec_model_request(
     params: JSONRPCParamsType,
     result_type: Type[BaseModelT],
     transport: JSONRPCTransport,
+    server_error_mapper: JSONRPCServerErrorMapper | None = None,
 ) -> BaseModelT:
     """
     Execute a JSON-RPC command to a sandbox environment expecting a model result.
@@ -110,6 +118,7 @@ async def exec_model_request(
       params (JSONRPCParamsType): The parameters for the JSON-RPC method.
       result_type (Type[BaseModelT]): The Pydantic model class to validate and parse the result.
       transport (JSONRPCTransport): The transport callable to use for the RPC communication.
+      server_error_mapper (JSONRPCServerErrorMapper): A callable to map server specific JSON-RPC errors to exceptions.
 
     Returns:
       BaseModelT: The parsed and validated result of the JSON-RPC call.
@@ -119,7 +128,12 @@ async def exec_model_request(
       ToolParsingError: If the JSON-RPC response contains a specific error code indicating a parsing error.
       ValueError: If the result cannot be validated against the provided model class.
     """
-    rpc_result = await _exec_request(method=method, params=params, transport=transport)
+    rpc_result = await _exec_request(
+        method=method,
+        params=params,
+        transport=transport,
+        server_error_mapper=server_error_mapper,
+    )
     return result_type.model_validate(rpc_result, strict=True)
 
 
@@ -161,6 +175,7 @@ async def _exec_request(
     method: str,
     params: JSONRPCParamsType,
     transport: JSONRPCTransport,
+    server_error_mapper: JSONRPCServerErrorMapper | None = None,
 ) -> object:
     """Execute a request using the provided transport mechanism."""
     return parse_json_rpc_response(
@@ -171,6 +186,7 @@ async def _exec_request(
         ),
         method,
         params,
+        server_error_mapper,
     )
 
 
@@ -178,15 +194,16 @@ def parse_json_rpc_response(
     response_str: str,
     method: str,
     params: JSONRPCParamsType,
+    server_error_mapper: JSONRPCServerErrorMapper | None = None,
 ) -> object:
     """Validates the JSON RPC response and returns the result or raises a proper Inspect error."""
     match JSONRPCResponse.model_validate_json(response_str).root:
         case JSONRPCSuccessResponse(result=rpc_result):
             return rpc_result
-        case JSONRPCErrorResponse(
-            error=JSONRPCError(code=code, message=message, data=_)
-        ):
-            raise exception_for_rpc_response_error(code, message, method, params)
+        case JSONRPCErrorResponse(error=JSONRPCError(code=code, message=message)):
+            raise exception_for_rpc_response_error(
+                code, message, method, params, server_error_mapper
+            )
         case _:
             raise ValueError(
                 f"Unexpected JSON RPC response to request {_rpc_call_description(method, params)}: {response_str}"
@@ -220,16 +237,17 @@ def exception_for_rpc_response_error(
             if server_error_mapper
             else ToolError(message)
         )
+    elif code == -32602:  # (Invalid params)
+        # Even though the Inspect side does validation, it can't possibly be
+        # complete - especially for tools that have dynamic action dependant
+        # rules for optional/required params.
+        return ToolParsingError(message)
     elif code == -32603:
         return ToolError(message)
     else:
         # -32600 (Invalid Request)
         #   If we sent a bogus request, it's 100% a code bug.
         # -32601 (Method not found)
-        # -32602 (Invalid params)
-        #   These shouldn't be possible since Inspect did validation prior to
-        #   making the tool call. Because of that, these errors should not make
-        #   it back to the model, so choose RuntimeError.
         # -32700 (Parse error)
         #   shouldn't be seen in this flow since we're processing responses, and
         #   this is a request oriented error.
@@ -276,10 +294,20 @@ def create_json_rpc_request(
     is_notification: bool,
 ) -> str:
     return json.dumps(
-        {
-            "jsonrpc": "2.0",
-            "method": method,
-            **({"params": params} if params else {}),
-            **({"id": next(id_generator)} if not is_notification else {}),
-        }
+        remove_none_values(
+            {
+                "jsonrpc": "2.0",
+                "method": method,
+                **({"params": params} if params else {}),
+                **({"id": next(id_generator)} if not is_notification else {}),
+            }
+        )
     )
+
+
+def remove_none_values(obj: object) -> object:
+    if isinstance(obj, dict):
+        return {k: remove_none_values(v) for k, v in obj.items() if v is not None}
+    elif isinstance(obj, list):
+        return [remove_none_values(item) for item in obj if item is not None]
+    return obj

inspect_ai/tool/_mcp/_mcp.py

@@ -259,6 +259,7 @@ def create_server_sandbox(
     cwd: str | Path | None = None,
     env: dict[str, str] | None = None,
     sandbox: str | None = None,
+    timeout: int | None = None,
 ) -> MCPServer:
     # TODO: Confirm the lifetime concepts. By the time a request makes it to the
     # sandbox, it's going to need both a session id and a server "name".
@@ -272,6 +273,7 @@ def create_server_sandbox(
                 env=env,
             ),
             sandbox_name=sandbox,
+            timeout=timeout,
         ),
         name=name,
         events=False,

inspect_ai/tool/_mcp/_sandbox.py

@@ -11,7 +11,7 @@ from inspect_ai.tool._tool_support_helpers import (
     exec_model_request,
     exec_notification,
     exec_scalar_request,
-    tool_container_sandbox,
+    tool_support_sandbox,
 )
 
 from ._context import MCPServerContext
@@ -28,8 +28,10 @@ async def sandbox_client(  # type: ignore
     *,
     sandbox_name: str | None = None,
     errlog: TextIO = sys.stderr,
+    timeout: int | None = None,  # default 180 seconds
 ) -> MCPServerContext:  # type: ignore
-    sandbox_environment = await tool_container_sandbox(
+    timeout = timeout or 180
+    (sandbox_environment, _) = await tool_support_sandbox(
         "mcp support", sandbox_name=sandbox_name
     )
 
@@ -49,6 +51,7 @@ async def sandbox_client(  # type: ignore
         method="mcp_launch_server",
         params={"server_params": server.model_dump()},
         result_type=int,
+        timeout=timeout,
     )
 
     async def stdout_reader() -> None:
@@ -72,6 +75,7 @@ async def sandbox_client(  # type: ignore
                                     "request": root.model_dump(),
                                 },
                                 result_type=JSONRPCMessage,
+                                timeout=timeout,
                             )
                         )
                     elif isinstance(root, JSONRPCNotification):
@@ -82,6 +86,7 @@ async def sandbox_client(  # type: ignore
                                 "session_id": session_id,
                                 "notification": root.model_dump(),
                             },
+                            timeout=timeout,
                         )
                     else:
                         assert False, f"Unexpected message type {message=}"
@@ -101,4 +106,5 @@ async def sandbox_client(  # type: ignore
                 method="mcp_kill_server",
                 params={"session_id": session_id},
                 result_type=type(None),
+                timeout=timeout,
             )

inspect_ai/tool/_mcp/server.py

@@ -73,6 +73,7 @@ def mcp_server_sandbox(
     cwd: str | Path | None = None,
     env: dict[str, str] | None = None,
     sandbox: str | None = None,
+    timeout: int | None = None,
 ) -> MCPServer:
     """MCP Server (Sandbox).
 
@@ -87,6 +88,7 @@ def mcp_server_sandbox(
             "SHELL", "TERM", and "USER" for Posix-based systems).
         cwd: The working directory to use when spawning the process.
         sandbox: The sandbox to use when spawning the process.
+        timeout: Timeout (in seconds) for command.
 
     Returns:
         McpClient: Client for MCP Server
@@ -94,7 +96,7 @@ def mcp_server_sandbox(
     verfify_mcp_package()
     from ._mcp import create_server_sandbox
 
-    return create_server_sandbox(command, args, cwd, env, sandbox)
+    return create_server_sandbox(command, args, cwd, env, sandbox, timeout)
 
 
 def verfify_mcp_package() -> None:

inspect_ai/tool/_tool_call.py

@@ -68,9 +68,12 @@ class ToolCallError:
         "permission",
         "file_not_found",
         "is_a_directory",
-        "output_limit",
+        "limit",
         "approval",
         "unknown",
+        # Retained for backward compatibility when loading logs created with an older
+        # version of inspect.
+        "output_limit",
     ]
     """Error type."""
 

inspect_ai/tool/_tool_support_helpers.py

@@ -7,13 +7,17 @@ It includes definitions for JSON-RPC request and response models, as well as fun
 from textwrap import dedent
 from typing import Type
 
+import semver
+
 from inspect_ai._util.error import PrerequisiteError
+from inspect_ai.tool._tool import ToolError
 from inspect_ai.util import sandbox_with
 from inspect_ai.util._sandbox.environment import SandboxEnvironment
 
 from ._json_rpc_helpers import (
     BaseModelT,
     JSONRPCParamsType,
+    JSONRPCServerErrorMapper,
     JSONRPCTransport,
     ScalarT,
     _rpc_call_description,
@@ -29,7 +33,7 @@ async def exec_scalar_request(
     method: str,
     params: JSONRPCParamsType,
     result_type: Type[ScalarT],
-    timeout: int | None = None,
+    timeout: int,
     user: str | None = None,
 ) -> ScalarT:
     return await scalar_request(
@@ -37,6 +41,7 @@ async def exec_scalar_request(
         params,
         result_type,
         transport=ToolSupportSandboxTransport(sandbox, timeout, user),
+        server_error_mapper=ToolSupportServerErrorMapper(),
     )
 
 
@@ -45,7 +50,7 @@ async def exec_model_request(
     method: str,
     params: JSONRPCParamsType,
     result_type: Type[BaseModelT],
-    timeout: int | None = None,
+    timeout: int,
     user: str | None = None,
 ) -> BaseModelT:
     return await model_request(
@@ -53,6 +58,7 @@ async def exec_model_request(
         params,
         result_type,
         transport=ToolSupportSandboxTransport(sandbox, timeout, user),
+        server_error_mapper=ToolSupportServerErrorMapper(),
     )
 
 
@@ -60,7 +66,7 @@ async def exec_notification(
     sandbox: SandboxEnvironment,
     method: str,
     params: JSONRPCParamsType,
-    timeout: int | None = None,
+    timeout: int,
     user: str | None = None,
 ) -> None:
     return await notification_helper(
@@ -68,19 +74,33 @@ async def exec_notification(
     )
 
 
+class ToolSupportServerErrorMapper(JSONRPCServerErrorMapper):
+    def __call__(
+        self, code: int, message: str, method: str, params: JSONRPCParamsType
+    ) -> Exception:
+        """Map `inspect-tool-support` defined custom codes to an exception."""
+        match code:
+            case -32099:  # This is a ToolException from the container
+                return ToolError(message)
+            case -32098:  # This is an unexpected exception inside the container
+                return RuntimeError(message)
+            case _:
+                return RuntimeError(message)
+
+
 class ToolSupportSandboxTransport(JSONRPCTransport):
     """
-    A transport callable that uses a sandbox for RPC communication.
+    A transport that uses a sandbox for RPC communication.
 
-    This class implements the TransportCallable protocol and encapsulates
-    the sandbox, timeout, and user parameters needed for sandbox-based
-    RPC communication.
+    This class implements the TransportCallable protocol and encapsulates the
+    sandbox, timeout, and user parameters needed for sandbox-based RPC
+    communication.
     """
 
     def __init__(
         self,
         sandbox: SandboxEnvironment,
-        timeout: int | None = None,
+        timeout: int,
         user: str | None = None,
     ):
         """
@@ -128,13 +148,32 @@ class ToolSupportSandboxTransport(JSONRPCTransport):
 
 SANDBOX_CLI = "inspect-tool-support"
 INSPECT_TOOL_SUPPORT_IMAGE_DOCKERHUB = "aisiuk/inspect-tool-support"
+FIRST_PUBLISHED_VERSION = semver.Version.parse("0.1.6")
+MIN_SUPPORTED_VERSION = FIRST_PUBLISHED_VERSION
+MIN_NON_DEPRECATED_VERSION = semver.Version.parse("1.0.0")
+
+
+async def _get_sandbox_tool_support_version(
+    sandbox: SandboxEnvironment,
+) -> semver.Version:
+    try:
+        return semver.Version.parse(
+            await exec_scalar_request(sandbox, "version", {}, str, 5)
+        )
+    except RuntimeError as rte:
+        if "-32601" in str(rte):
+            # The container doesn't even have a version method. The first version
+            # published was 0.1.6, so we'll have to assume it was that old.
+            return FIRST_PUBLISHED_VERSION
+        raise rte
 
 
-async def tool_container_sandbox(
+async def tool_support_sandbox(
     tool_name: str, *, sandbox_name: str | None = None
-) -> SandboxEnvironment:
+) -> tuple[SandboxEnvironment, semver.Version]:
     if sb := await sandbox_with(SANDBOX_CLI, True, name=sandbox_name):
-        return sb
+        current_version = await _get_sandbox_tool_support_version(sb)
+        return (sb, current_version)
 
     # This sort of programmatic sentence building will not cut it if we ever
     # support other languages.
@@ -160,7 +199,7 @@ async def tool_container_sandbox(
 
 
 def create_sandbox_transport(
-    sandbox: SandboxEnvironment, timeout: int | None = None, user: str | None = None
+    sandbox: SandboxEnvironment, timeout: int, user: str | None = None
 ) -> JSONRPCTransport:
     """
     Create a transport callable that uses a sandbox for RPC communication.