inspect-ai 0.3.72__py3-none-any.whl → 0.3.73__py3-none-any.whl

inspect_ai/tool/_tool_params.py

@@ -1,44 +1,14 @@
 from typing import (
-    Any,
     Literal,
-    Optional,
+    TypeAlias,
 )
 
 from pydantic import BaseModel, Field
 
-JSONType = Literal["string", "integer", "number", "boolean", "array", "object", "null"]
-"""Validate types within JSON schema."""
+from inspect_ai.util._json import JSONSchema
 
-
-class ToolParam(BaseModel):
-    """Description of tool parameter in JSON Schema format."""
-
-    type: JSONType | None = Field(default=None)
-    """JSON type of tool parameter."""
-
-    description: str | None = Field(default=None)
-    """Parameter description."""
-
-    default: Any = Field(default=None)
-    """Default value for parameter."""
-
-    enum: list[Any] | None = Field(default=None)
-    """Valid values for enum parameters."""
-
-    items: Optional["ToolParam"] = Field(default=None)
-    """Valid type for array parameters."""
-
-    properties: dict[str, "ToolParam"] | None = Field(default=None)
-    """Valid fields for object parametrs."""
-
-    additionalProperties: Optional["ToolParam"] | bool | None = Field(default=None)
-    """Are additional properties allowed?"""
-
-    anyOf: list["ToolParam"] | None = Field(default=None)
-    """Valid types for union parameters."""
-
-    required: list[str] | None = Field(default=None)
-    """Required fields for object parameters."""
+ToolParam: TypeAlias = JSONSchema
+"""Description of tool parameter in JSON Schema format."""
 
 
 class ToolParams(BaseModel):

inspect_ai/tool/_tools/_web_search.py

@@ -1,7 +1,7 @@
-import asyncio
 import os
 from typing import Literal, Protocol, runtime_checkable
 
+import anyio
 import httpx
 from bs4 import BeautifulSoup, NavigableString
 from tenacity import (
@@ -13,7 +13,7 @@ from tenacity import (
 )
 
 from inspect_ai._util.error import PrerequisiteError
-from inspect_ai._util.retry import httpx_should_retry, log_retry_attempt
+from inspect_ai._util.httpx import httpx_should_retry, log_httpx_retry_attempt
 from inspect_ai.util._concurrency import concurrency
 
 from .._tool import Tool, ToolResult, tool
@@ -25,6 +25,17 @@ Page Content: {text}
 """
 
 
+class SearchLink:
+    def __init__(self, url: str, snippet: str) -> None:
+        self.url = url
+        self.snippet = snippet
+
+
+@runtime_checkable
+class SearchProvider(Protocol):
+    async def __call__(self, query: str, start_idx: int) -> list[SearchLink]: ...
+
+
 @tool
 def web_search(
     provider: Literal["google"] = "google",
@@ -84,16 +95,22 @@ def web_search(
             async with concurrency(f"{provider}_web_search", max_connections):
                 links = await search_provider(query, start_idx=search_calls * 10)
 
-            # Extract and summarize each page individually
-            pages = await asyncio.gather(
-                *[page_if_relevant(link.url, query, model, client) for link in links],
-                return_exceptions=True,
-            )
-            for page, link in zip(pages, links):
-                if page and not isinstance(page, BaseException):
-                    page_contents.append(page)
-                    urls.append(link.url)
-                    snippets.append(link.snippet)
+            async with anyio.create_task_group() as tg:
+
+                async def process_link(link: SearchLink) -> None:
+                    try:
+                        page = await page_if_relevant(link.url, query, model, client)
+                        if page:
+                            page_contents.append(page)
+                            urls.append(link.url)
+                            snippets.append(link.snippet)
+                    # exceptions fetching pages are very common!
+                    except Exception:
+                        pass
+
+                for lk in links:
+                    tg.start_soon(process_link, lk)
+
             search_calls += 1
 
         all_page_contents = "\n\n".join(page_contents)
@@ -168,17 +185,6 @@ async def page_if_relevant(
         return None
 
 
-class SearchLink:
-    def __init__(self, url: str, snippet: str) -> None:
-        self.url = url
-        self.snippet = snippet
-
-
-@runtime_checkable
-class SearchProvider(Protocol):
-    async def __call__(self, query: str, start_idx: int) -> list[SearchLink]: ...
-
-
 def google_search_provider(client: httpx.AsyncClient) -> SearchProvider:
     google_api_key = os.environ.get("GOOGLE_CSE_API_KEY", None)
     google_cse_id = os.environ.get("GOOGLE_CSE_ID", None)
@@ -204,7 +210,7 @@ def google_search_provider(client: httpx.AsyncClient) -> SearchProvider:
             wait=wait_exponential_jitter(),
             stop=stop_after_attempt(5) | stop_after_delay(60),
             retry=retry_if_exception(httpx_should_retry),
-            before_sleep=log_retry_attempt(search_url),
+            before_sleep=log_httpx_retry_attempt(search_url),
         )
         async def execute_search() -> httpx.Response:
             return await client.get(search_url)

inspect_ai/util/__init__.py

@@ -3,6 +3,7 @@ from inspect_ai._util.trace import trace_action, trace_message
 from ._concurrency import concurrency
 from ._console import input_screen
 from ._display import DisplayType, display_counter, display_type
+from ._json import JSONSchema, JSONType, json_schema
 from ._panel import InputPanel, input_panel
 from ._resource import resource
 from ._sandbox import (
@@ -36,6 +37,9 @@ __all__ = [
     "InputPanel",
     "input_panel",
     "input_screen",
+    "JSONType",
+    "JSONSchema",
+    "json_schema",
     "OutputLimitExceededError",
     "resource",
     "subprocess",

inspect_ai/util/_concurrency.py

@@ -1,9 +1,10 @@
-import asyncio
 import contextlib
 import time
 from dataclasses import dataclass
 from typing import AsyncIterator
 
+import anyio
+
 from inspect_ai._util.working import report_sample_waiting_time
 
 
@@ -45,9 +46,7 @@ async def concurrency(
     # do we have an existing semaphore? if not create one and store it
     semaphore = _concurrency_semaphores.get(key, None)
     if semaphore is None:
-        semaphore = ConcurencySempahore(
-            name, concurrency, asyncio.Semaphore(concurrency)
-        )
+        semaphore = ConcurencySempahore(name, concurrency, anyio.Semaphore(concurrency))
         _concurrency_semaphores[key] = semaphore
 
     # wait and yield to protected code
@@ -60,7 +59,7 @@ async def concurrency(
 def concurrency_status() -> dict[str, tuple[int, int]]:
     status: dict[str, tuple[int, int]] = {}
     for c in _concurrency_semaphores.values():
-        status[c.name] = (c.concurrency - c.semaphore._value, c.concurrency)
+        status[c.name] = (c.concurrency - c.semaphore.value, c.concurrency)
     return status
 
 
@@ -72,7 +71,7 @@ def init_concurrency() -> None:
 class ConcurencySempahore:
     name: str
     concurrency: int
-    semaphore: asyncio.Semaphore
+    semaphore: anyio.Semaphore
 
 
 _concurrency_semaphores: dict[str, ConcurencySempahore] = {}

inspect_ai/util/_display.py

@@ -2,6 +2,7 @@ import os
 from logging import getLogger
 from typing import Literal
 
+from inspect_ai._util._async import configured_async_backend
 from inspect_ai._util.constants import DEFAULT_DISPLAY
 from inspect_ai._util.thread import is_main_thread
 
@@ -20,6 +21,11 @@ def init_display_type(display: str | None = None) -> DisplayType:
         display or os.environ.get("INSPECT_DISPLAY", DEFAULT_DISPLAY).lower().strip()
     )
 
+    # if trio is configured as the backend then throttle down to "rich"
+    # (as textual uses asyncio directly so is not compatible with trio)
+    if configured_async_backend() == "trio" and display == "full":
+        display = "rich"
+
     # if we are on a background thread then throttle down to "plain"
     # ("full" requires textual which cannot run in a background thread
     # b/c it calls the Python signal function; "rich" assumes exclusive

inspect_ai/util/_json.py

@@ -0,0 +1,170 @@
+import types
+import typing
+from dataclasses import is_dataclass
+from typing import (
+    Any,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    Type,
+    Union,
+    get_args,
+    get_origin,
+    get_type_hints,
+    is_typeddict,
+)
+
+from pydantic import BaseModel, Field
+
+JSONType = Literal["string", "integer", "number", "boolean", "array", "object", "null"]
+"""Valid types within JSON schema."""
+
+
+class JSONSchema(BaseModel):
+    """JSON Schema for type."""
+
+    type: JSONType | None = Field(default=None)
+    """JSON type of tool parameter."""
+
+    description: str | None = Field(default=None)
+    """Parameter description."""
+
+    default: Any = Field(default=None)
+    """Default value for parameter."""
+
+    enum: list[Any] | None = Field(default=None)
+    """Valid values for enum parameters."""
+
+    items: Optional["JSONSchema"] = Field(default=None)
+    """Valid type for array parameters."""
+
+    properties: dict[str, "JSONSchema"] | None = Field(default=None)
+    """Valid fields for object parametrs."""
+
+    additionalProperties: Optional["JSONSchema"] | bool | None = Field(default=None)
+    """Are additional properties allowed?"""
+
+    anyOf: list["JSONSchema"] | None = Field(default=None)
+    """Valid types for union parameters."""
+
+    required: list[str] | None = Field(default=None)
+    """Required fields for object parameters."""
+
+
+def json_schema(t: Type[Any]) -> JSONSchema:
+    """Provide a JSON Schema for the specified type.
+
+    Schemas can be automatically inferred for a wide variety of
+    Python class types including Pydantic BaseModel, dataclasses,
+    and typed dicts.
+
+    Args:
+        t: Python type
+
+    Returns:
+        JSON Schema for type.
+    """
+    origin = get_origin(t)
+    args = get_args(t)
+
+    if origin is None:
+        if t is int:
+            return JSONSchema(type="integer")
+        elif t is float:
+            return JSONSchema(type="number")
+        elif t is str:
+            return JSONSchema(type="string")
+        elif t is bool:
+            return JSONSchema(type="boolean")
+        elif t is list:
+            return JSONSchema(type="array", items=JSONSchema())
+        elif t is dict:
+            return JSONSchema(type="object", additionalProperties=JSONSchema())
+        elif (
+            is_dataclass(t)
+            or is_typeddict(t)
+            or (isinstance(t, type) and issubclass(t, BaseModel))
+        ):
+            return cls_json_schema(t)
+        elif t is type(None):
+            return JSONSchema(type="null")
+        else:
+            return JSONSchema()
+    elif origin is list or origin is List or origin is tuple or origin is Tuple:
+        return JSONSchema(
+            type="array", items=json_schema(args[0]) if args else JSONSchema()
+        )
+    elif origin is dict or origin is Dict:
+        return JSONSchema(
+            type="object",
+            additionalProperties=json_schema(args[1])
+            if len(args) > 1
+            else JSONSchema(),
+        )
+    elif origin is Union or origin is types.UnionType:
+        return JSONSchema(anyOf=[json_schema(arg) for arg in args])
+    elif origin is Optional:
+        return JSONSchema(
+            anyOf=[json_schema(arg) for arg in args] + [JSONSchema(type="null")]
+        )
+    elif origin is typing.Literal:
+        return JSONSchema(enum=list(args))
+
+    return JSONSchema()  # Default case if we can't determine the type
+
+
+def cls_json_schema(cls: Type[Any]) -> JSONSchema:
+    properties: Dict[str, JSONSchema] = {}
+    required: List[str] = []
+
+    if is_dataclass(cls):
+        fields = cls.__dataclass_fields__  # type: ignore
+        for name, field in fields.items():
+            properties[name] = json_schema(field.type)  # type: ignore
+            if field.default == field.default_factory:
+                required.append(name)
+    elif isinstance(cls, type) and issubclass(cls, BaseModel):
+        schema = cls.model_json_schema()
+        for name, prop in schema.get("properties", {}).items():
+            properties[name] = JSONSchema(**prop)
+        required = schema.get("required", [])
+    elif is_typeddict(cls):
+        annotations = get_type_hints(cls)
+        for name, type_hint in annotations.items():
+            properties[name] = json_schema(type_hint)
+            if name in cls.__required_keys__:
+                required.append(name)
+
+    return JSONSchema(
+        type="object",
+        properties=properties,
+        required=required if required else None,
+        additionalProperties=False,
+    )
+
+
+def python_type_to_json_type(python_type: str | None) -> JSONType:
+    match python_type:
+        case "str":
+            return "string"
+        case "int":
+            return "integer"
+        case "float":
+            return "number"
+        case "bool":
+            return "boolean"
+        case "list":
+            return "array"
+        case "dict":
+            return "object"
+        case "None":
+            return "null"
+        # treat 'unknown' as string as anything can be converted to string
+        case None:
+            return "string"
+        case _:
+            raise ValueError(
+                f"Unsupported type: {python_type} for Python to JSON conversion."
+            )

inspect_ai/util/_sandbox/docker/cleanup.py

@@ -1,12 +1,14 @@
-import asyncio
 from contextvars import ContextVar
 from pathlib import Path
 from typing import Awaitable, Callable, Set
 
+import anyio
 from rich import box, print
 from rich.panel import Panel
 from rich.table import Table
 
+from inspect_ai._util._async import coro_print_exceptions
+
 from .compose import compose_down, compose_ls, compose_ps
 from .config import is_auto_compose_file, safe_cleanup_auto_compose
 from .util import ComposeProject
@@ -94,13 +96,15 @@ async def cleanup_projects(
     )
 
     # cleanup all of the projects in parallel
-    tasks = [cleanup_fn(project, False) for project in projects]
-    results = await asyncio.gather(*tasks, return_exceptions=True)
-
-    # report errors
-    for result in results:
-        if result is not None:
-            print(f"Error cleaning up Docker environment: {result}")
+    async with anyio.create_task_group() as tg:
+        for project in projects:
+            tg.start_soon(
+                coro_print_exceptions,
+                "cleaning up Docker environment",
+                cleanup_fn,
+                project,
+                False,
+            )
 
 
 async def cli_cleanup(project_name: str | None) -> None:
@@ -141,7 +145,7 @@ def auto_compose_files() -> Set[str]:
 
 
 _running_projects: ContextVar[list[ComposeProject]] = ContextVar(
-    "docker_running_projects"
+    "docker_running_projects", default=[]
 )
 
 _auto_compose_files: ContextVar[Set[str]] = ContextVar("docker_auto_compose_files")

inspect_ai/util/_sandbox/docker/docker.py

@@ -293,6 +293,9 @@ class DockerSandboxEnvironment(SandboxEnvironment):
 
     @override
     async def write_file(self, file: str, contents: str | bytes) -> None:
+        # defualt timeout for write_file operations
+        TIMEOUT = 180
+
         # resolve relative file paths
         file = self.container_file(file)
 
@@ -309,6 +312,7 @@ class DockerSandboxEnvironment(SandboxEnvironment):
             result = await self.exec(
                 ["sh", "-e", "-c", 'tee -- "$1"', "write_file_script", file],
                 input=contents,
+                timeout=TIMEOUT,
             )
         else:
             base64_contents = base64.b64encode(contents).decode("US-ASCII")
@@ -322,6 +326,7 @@ class DockerSandboxEnvironment(SandboxEnvironment):
                     file,
                 ],
                 input=base64_contents,
+                timeout=TIMEOUT,
             )
         if result.returncode != 0:
             if "permission denied" in result.stderr.casefold():

inspect_ai/util/_sandbox/environment.py

@@ -3,18 +3,19 @@ from __future__ import annotations
 import abc
 from dataclasses import dataclass, field
 from typing import (
+    Annotated,
     Any,
     Awaitable,
     Callable,
     Literal,
-    NamedTuple,
     Type,
     TypeVar,
     Union,
+    cast,
     overload,
 )
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, model_validator
 
 from .._subprocess import ExecResult
 
@@ -38,6 +39,7 @@ SampleCleanup = Callable[
     ],
     Awaitable[None],
 ]
+ConfigDeserialize = Callable[[dict[str, Any]], BaseModel]
 
 
 class HostMapping(BaseModel):
@@ -211,11 +213,6 @@ class SandboxEnvironment(abc.ABC):
                 f"Expected instance of {sandbox_cls.__name__}, got {type(self).__name__}"
             )
 
-    @classmethod
-    def config_files(cls) -> list[str]:
-        """Standard config files for this provider (used for automatic discovery)"""
-        return []
-
     @classmethod
     def default_concurrency(cls) -> int | None:
         """Default max_sandboxes for this provider (`None` means no maximum)"""
@@ -296,6 +293,30 @@ class SandboxEnvironment(abc.ABC):
         """
         pass
 
+    @classmethod
+    def config_files(cls) -> list[str]:
+        """Standard config files for this provider (used for automatic discovery)"""
+        return []
+
+    @classmethod
+    def config_deserialize(cls, config: dict[str, Any]) -> BaseModel:
+        """Deserialize a sandbox-specific configuration model from a dict.
+
+        Override this method if you support a custom configuration model.
+
+        A basic implementation would be: `return MySandboxEnvironmentConfig(**config)`
+
+        Args:
+          config: Configuration dictionary produced by serializing the configuration
+            model.
+
+        Returns:
+          The sandbox-specific configuration model.
+        """
+        raise NotImplementedError(
+            "The SandboxEnvironment provider has not implemented config_deserialize."
+        )
+
 
 @dataclass
 class SandboxEnvironments:
@@ -311,15 +332,30 @@ class SandboxEnvironments:
     """
 
 
-class SandboxEnvironmentSpec(NamedTuple):
+class SandboxEnvironmentSpec(BaseModel, frozen=True):
     """Specification of a SandboxEnvironment."""
 
     type: str
     """Sandbox type (e.g. 'local', 'docker')"""
 
-    config: SandboxEnvironmentConfigType | None = None
+    # Any is used to prevent Pydantic from trying to initialise a BaseModel.
+    config: Annotated[Any, "BaseModel, str or None"] = None
     """Sandbox configuration (filename or config object)."""
 
+    def __init__(self, type: str, config: BaseModel | str | None = None):
+        super().__init__(type=type, config=config)
+
+    @model_validator(mode="before")
+    @classmethod
+    def load_config_model(cls, data: dict[str, Any]) -> dict[str, Any]:
+        type = data["type"]
+        config = data.get("config")
+        # Pydantic won't know what concrete type to instantiate for config, so
+        # ask the relevant sandbox environment to deserialize it.
+        if isinstance(config, dict) and len(config) > 0:
+            data["config"] = deserialize_sandbox_specific_config(type, config)
+        return data
+
 
 SandboxEnvironmentConfigType = BaseModel | str
 
@@ -343,3 +379,14 @@ def resolve_sandbox_environment(
         return SandboxEnvironmentSpec(sandbox[0], sandbox[1])
     else:
         return None
+
+
+def deserialize_sandbox_specific_config(type: str, config: dict[str, Any]) -> BaseModel:
+    # Avoid circular import
+    from inspect_ai.util._sandbox.registry import registry_find_sandboxenv
+
+    sandboxenv_type = registry_find_sandboxenv(type)
+    config_deserialize = cast(
+        ConfigDeserialize, getattr(sandboxenv_type, "config_deserialize")
+    )
+    return config_deserialize(config)

inspect_ai/util/_sandbox/service.py

@@ -1,4 +1,3 @@
-import asyncio
 import json
 from logging import getLogger
 from pathlib import PurePosixPath
@@ -9,8 +8,10 @@ from typing import (
     cast,
 )
 
+import anyio
 from pydantic import JsonValue
 
+from inspect_ai._util._async import coro_log_exceptions
 from inspect_ai.util._subprocess import ExecResult
 
 from .environment import SandboxEnvironment
@@ -59,7 +60,7 @@ async def sandbox_service(
 
     # wait for and process methods
     while not until():
-        await asyncio.sleep(POLLING_INTERVAL)
+        await anyio.sleep(POLLING_INTERVAL)
         await service.handle_requests()
 
 
@@ -141,9 +142,15 @@ class SandboxService:
         if result.success:
             request_files = result.stdout.strip().splitlines()
             if request_files:
-                await asyncio.gather(
-                    *[self._handle_request(file) for file in request_files]
-                )
+                async with anyio.create_task_group() as tg:
+                    for file in request_files:
+                        tg.start_soon(
+                            coro_log_exceptions,
+                            logger,
+                            "handling sandbox service request",
+                            self._handle_request,
+                            file,
+                        )
 
     async def _handle_request(self, request_file: str) -> None:
         # read request