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

inspect_ai/model/_providers/util/__init__.py

@@ -5,7 +5,7 @@ from .chatapi import (
     ChatAPIMessage,
     chat_api_input,
     chat_api_request,
-    is_chat_api_rate_limit,
+    should_retry_chat_api_error,
 )
 from .hf_handler import HFHandler
 from .llama31 import Llama31Handler
@@ -19,7 +19,7 @@ __all__ = [
     "as_stop_reason",
     "chat_api_request",
     "chat_api_input",
-    "is_chat_api_rate_limit",
+    "should_retry_chat_api_error",
     "model_base_url",
     "parse_tool_call",
     "tool_parse_error_message",

inspect_ai/model/_providers/util/chatapi.py

@@ -7,17 +7,15 @@ from tenacity import (
     retry,
     retry_if_exception,
     stop_after_attempt,
-    stop_after_delay,
     wait_exponential_jitter,
 )
 
-from inspect_ai._util.constants import DEFAULT_MAX_RETRIES
-from inspect_ai._util.retry import httpx_should_retry, log_retry_attempt
+from inspect_ai._util.http import is_retryable_http_status
+from inspect_ai._util.httpx import httpx_should_retry, log_httpx_retry_attempt
 from inspect_ai.model._chat_message import ChatMessageAssistant, ChatMessageTool
 from inspect_ai.tool._tool_info import ToolInfo
 
 from ..._chat_message import ChatMessage
-from ..._generate_config import GenerateConfig
 
 logger = getLogger(__name__)
 
@@ -75,21 +73,13 @@ async def chat_api_request(
     url: str,
     headers: dict[str, Any],
     json: Any,
-    config: GenerateConfig,
 ) -> Any:
-    # provide default max_retries
-    max_retries = config.max_retries if config.max_retries else DEFAULT_MAX_RETRIES
-
     # define call w/ retry policy
     @retry(
         wait=wait_exponential_jitter(),
-        stop=(
-            (stop_after_attempt(max_retries) | stop_after_delay(config.timeout))
-            if config.timeout
-            else stop_after_attempt(max_retries)
-        ),
+        stop=(stop_after_attempt(2)),
         retry=retry_if_exception(httpx_should_retry),
-        before_sleep=log_retry_attempt(model_name),
+        before_sleep=log_httpx_retry_attempt(model_name),
     )
     async def call_api() -> Any:
         response = await client.post(url=url, headers=headers, json=json)
@@ -104,14 +94,11 @@ async def chat_api_request(
 # checking for rate limit errors needs to punch through the RetryError and
 # look at its `__cause__`. we've observed Cloudflare giving transient 500
 # status as well as a ReadTimeout, so we count these as rate limit errors
-def is_chat_api_rate_limit(ex: BaseException) -> bool:
+def should_retry_chat_api_error(ex: BaseException) -> bool:
     return isinstance(ex, RetryError) and (
         (
             isinstance(ex.__cause__, httpx.HTTPStatusError)
-            and (
-                ex.__cause__.response.status_code == 429
-                or ex.__cause__.response.status_code == 500
-            )
+            and is_retryable_http_status(ex.__cause__.response.status_code)
         )
         or isinstance(ex.__cause__, httpx.ReadTimeout)
     )

inspect_ai/model/_providers/util/hooks.py

@@ -0,0 +1,165 @@
+import re
+import time
+from logging import getLogger
+from typing import Any, Mapping, NamedTuple, cast
+
+import httpx
+from shortuuid import uuid
+
+from inspect_ai._util.constants import HTTP
+from inspect_ai._util.retry import report_http_retry
+
+logger = getLogger(__name__)
+
+
+class RequestInfo(NamedTuple):
+    attempts: int
+    last_request: float
+
+
+class HttpHooks:
+    """Class which hooks various HTTP clients for improved tracking/logging.
+
+    A special header is injected into requests which is then read from
+    a request event hook -- this creates a record of when the request
+    started. Note that with retries a single request_id could be started
+    several times; our request hook makes sure we always track the time of
+    the last request.
+
+    There is an 'end_request()' method which gets the total request time
+    for a request_id and then purges the request_id from our tracking (so
+    the dict doesn't grow unbounded)
+
+    Additionally, an http response hook is installed and used for logging
+    requests for the 'http' log-level
+    """
+
+    REQUEST_ID_HEADER = "x-irid"
+
+    def __init__(self) -> None:
+        # track request start times
+        self._requests: dict[str, RequestInfo] = {}
+
+    def start_request(self) -> str:
+        request_id = uuid()
+        self._requests[request_id] = RequestInfo(0, time.monotonic())
+        return request_id
+
+    def end_request(self, request_id: str) -> float:
+        # read the request info (if available) and purge from dict
+        request_info = self._requests.pop(request_id, None)
+        if request_info is None:
+            raise RuntimeError(f"request_id not registered: {request_id}")
+
+        # return elapsed time
+        return time.monotonic() - request_info.last_request
+
+    def update_request_time(self, request_id: str) -> None:
+        request_info = self._requests.get(request_id, None)
+        if not request_info:
+            raise RuntimeError(f"No request registered for request_id: {request_id}")
+
+        # update the attempts and last request time
+        request_info = RequestInfo(request_info.attempts + 1, time.monotonic())
+        self._requests[request_id] = request_info
+
+        # trace a retry if this is attempt > 1
+        if request_info.attempts > 1:
+            report_http_retry()
+
+
+class ConverseHooks(HttpHooks):
+    def __init__(self, session: Any) -> None:
+        from aiobotocore.session import AioSession
+
+        super().__init__()
+
+        # register hooks
+        session = cast(AioSession, session._session)
+        session.register(
+            "before-send.bedrock-runtime.Converse", self.converse_before_send
+        )
+        session.register(
+            "after-call.bedrock-runtime.Converse", self.converse_after_call
+        )
+
+    def converse_before_send(self, **kwargs: Any) -> None:
+        user_agent = kwargs["request"].headers["User-Agent"].decode()
+        match = re.search(rf"{self.USER_AGENT_PREFIX}(\w+)", user_agent)
+        if match:
+            request_id = match.group(1)
+            self.update_request_time(request_id)
+
+    def converse_after_call(self, http_response: Any, **kwargs: Any) -> None:
+        from botocore.awsrequest import AWSResponse
+
+        response = cast(AWSResponse, http_response)
+        logger.log(HTTP, f"POST {response.url} - {response.status_code}")
+
+    def user_agent_extra(self, request_id: str) -> str:
+        return f"{self.USER_AGENT_PREFIX}{request_id}"
+
+    USER_AGENT_PREFIX = "ins/rid#"
+
+
+class HttpxHooks(HttpHooks):
+    def __init__(self, client: httpx.AsyncClient):
+        super().__init__()
+
+        # install hooks
+        client.event_hooks["request"].append(self.request_hook)
+        client.event_hooks["response"].append(self.response_hook)
+
+    async def request_hook(self, request: httpx.Request) -> None:
+        # update the last request time for this request id (as there could be retries)
+        request_id = request.headers.get(self.REQUEST_ID_HEADER, None)
+        if request_id:
+            self.update_request_time(request_id)
+
+    async def response_hook(self, response: httpx.Response) -> None:
+        message = f'{response.request.method} {response.request.url} "{response.http_version} {response.status_code} {response.reason_phrase}" '
+        logger.log(HTTP, message)
+
+
+def urllib3_hooks() -> HttpHooks:
+    import urllib3
+    from urllib3.connectionpool import HTTPConnectionPool
+    from urllib3.response import BaseHTTPResponse
+
+    class Urllib3Hooks(HttpHooks):
+        def request_hook(self, headers: Mapping[str, str]) -> None:
+            # update the last request time for this request id (as there could be retries)
+            request_id = headers.get(self.REQUEST_ID_HEADER, None)
+            if request_id:
+                self.update_request_time(request_id)
+
+        def response_hook(
+            self, method: str, url: str, response: BaseHTTPResponse
+        ) -> None:
+            message = f'{method} {url} "{response.version_string} {response.status} {response.reason}" '
+            logger.log(HTTP, message)
+
+    global _urlilb3_hooks
+    if _urlilb3_hooks is None:
+        # one time patch of urlopen
+        urlilb3_hooks = Urllib3Hooks()
+        original_urlopen = urllib3.connectionpool.HTTPConnectionPool.urlopen
+
+        def patched_urlopen(
+            self: HTTPConnectionPool, method: str, url: str, **kwargs: Any
+        ) -> BaseHTTPResponse:
+            headers = kwargs.get("headers", {})
+            urlilb3_hooks.request_hook(headers)
+            response = original_urlopen(self, method, url, **kwargs)
+            urlilb3_hooks.response_hook(method, f"{self.host}{url}", response)
+            return response
+
+        urllib3.connectionpool.HTTPConnectionPool.urlopen = patched_urlopen  # type: ignore[assignment,method-assign]
+
+        # assign to global hooks instance
+        _urlilb3_hooks = urlilb3_hooks
+
+    return _urlilb3_hooks
+
+
+_urlilb3_hooks: HttpHooks | None = None

inspect_ai/model/_providers/vertex.py

@@ -4,7 +4,13 @@ from copy import copy
 from typing import Any, cast
 
 import vertexai  # type: ignore
-from google.api_core.exceptions import TooManyRequests
+from google.api_core.exceptions import (
+    Aborted,
+    ClientError,
+    DeadlineExceeded,
+    ServiceUnavailable,
+)
+from google.api_core.retry import if_transient_error
 from google.protobuf.json_format import MessageToDict
 from pydantic import JsonValue
 from typing_extensions import override
@@ -31,6 +37,7 @@ from inspect_ai._util.content import (
     ContentText,
     ContentVideo,
 )
+from inspect_ai._util.http import is_retryable_http_status
 from inspect_ai._util.images import file_as_data
 from inspect_ai.tool import ToolCall, ToolChoice, ToolInfo
 
@@ -169,8 +176,18 @@ class VertexAPI(ModelAPI):
         return output, call
 
     @override
-    def is_rate_limit(self, ex: BaseException) -> bool:
-        return isinstance(ex, TooManyRequests)
+    def should_retry(self, ex: Exception) -> bool:
+        # google API-specific errors
+        if isinstance(ex, Aborted | DeadlineExceeded | ServiceUnavailable):
+            return True
+        # standard HTTP errors
+        elif isinstance(ex, ClientError) and ex.code is not None:
+            return is_retryable_http_status(ex.code)
+        # additional errors flagged by google as transient
+        elif isinstance(ex, Exception):
+            return if_transient_error(ex)
+        else:
+            return False
 
     @override
     def connection_key(self) -> str:

inspect_ai/model/_providers/vllm.py

@@ -1,13 +1,15 @@
-import asyncio
+import concurrent.futures
 import functools
 import gc
 import os
 import time
+from concurrent.futures import Future
 from dataclasses import dataclass
 from queue import Empty, Queue
 from threading import Thread
 from typing import Any, cast
 
+import anyio
 from typing_extensions import override
 from vllm import LLM, CompletionOutput, RequestOutput, SamplingParams  # type: ignore
 
@@ -280,8 +282,7 @@ class VLLMAPI(ModelAPI):
 @dataclass
 class _QueueItem:
     input: GenerateInput
-    future: asyncio.Future[list[GenerateOutput]]
-    loop: asyncio.AbstractEventLoop
+    future: Future[list[GenerateOutput]]
 
 
 batch_thread: Thread | None = None
@@ -297,15 +298,16 @@ async def batched_generate(input: GenerateInput) -> list[GenerateOutput]:
         batch_thread.start()
 
     # enqueue the job
-    loop = asyncio.get_event_loop()
-    future: asyncio.Future[list[GenerateOutput]] = loop.create_future()
-    batch_queue.put(_QueueItem(input=input, future=future, loop=loop))
+    future = Future[list[GenerateOutput]]()
+    batch_queue.put(_QueueItem(input=input, future=future))
 
-    # await the job
-    await future
-
-    # return it
-    return future.result()
+    # await the future
+    while True:
+        try:
+            return future.result(timeout=0.01)
+        except concurrent.futures.TimeoutError:
+            pass
+        await anyio.sleep(1)
 
 
 def string_to_bytes(string: str) -> list[int]:
@@ -397,13 +399,12 @@ def post_process_outputs(
 def process_batches() -> None:
     while True:
         # drain the queue (wait until no new messages have shown up for 2 seconds)
-        inputs: list[tuple[GenerateInput, asyncio.Future[list[GenerateOutput]]]] = []
+        inputs: list[tuple[GenerateInput, Future[list[GenerateOutput]]]] = []
         while True:
             try:
                 input = batch_queue.get(
                     timeout=2
                 )  # wait 2 seconds max TODO: what's optimal wait time?
-                loop = input.loop
                 inputs.append((input.input, input.future))
                 if len(inputs) >= input.input.batch_size:
                     # max batch size reached
@@ -429,14 +430,10 @@ def process_batches() -> None:
             for i, output in enumerate(outputs):
                 future = inputs[i][1]
 
-                # asyncio futures are not thread safe, so we need to pass the event loop
-                # down to this point, so we can mark the future as done in a thread safe manner.
-                # see: https://docs.python.org/3/library/asyncio-dev.html#concurrency-and-multithreading
-                loop.call_soon_threadsafe(
-                    future.set_result,
+                future.set_result(
                     post_process_outputs(output, num_top_logprobs, total_time),
                 )
 
         except Exception as e:
             for _, future in inputs:
-                loop.call_soon_threadsafe(future.set_exception, e)
+                future.set_exception(e)

inspect_ai/scorer/_multi.py

@@ -1,5 +1,6 @@
-import asyncio
+import functools
 
+from inspect_ai._util._async import tg_collect
 from inspect_ai.scorer._reducer.registry import create_reducers
 from inspect_ai.solver._task_state import TaskState
 
@@ -19,7 +20,9 @@ def multi_scorer(scorers: list[Scorer], reducer: str | ScoreReducer) -> Scorer:
     reducer = create_reducers(reducer)[0]
 
     async def score(state: TaskState, target: Target) -> Score:
-        scores = await asyncio.gather(*[_scorer(state, target) for _scorer in scorers])
+        scores = await tg_collect(
+            [functools.partial(_scorer, state, target) for _scorer in scorers]
+        )
         return reducer(scores)
 
     return score

inspect_ai/solver/_bridge/patch.py

@@ -11,11 +11,12 @@ from openai._types import ResponseT
 from openai.types.chat import (
     ChatCompletion,
     ChatCompletionMessageParam,
+    ChatCompletionToolChoiceOptionParam,
     ChatCompletionToolParam,
 )
 from shortuuid import uuid
 
-from inspect_ai.model._generate_config import GenerateConfig
+from inspect_ai.model._generate_config import GenerateConfig, ResponseSchema
 from inspect_ai.model._model import get_model
 from inspect_ai.model._openai import (
     chat_messages_from_openai,
@@ -23,8 +24,10 @@ from inspect_ai.model._openai import (
     openai_completion_usage,
 )
 from inspect_ai.solver._task_state import sample_state
+from inspect_ai.tool._tool_choice import ToolChoice, ToolFunction
 from inspect_ai.tool._tool_info import ToolInfo
 from inspect_ai.tool._tool_params import ToolParams
+from inspect_ai.util._json import JSONSchema
 
 
 @contextlib.asynccontextmanager
@@ -113,6 +116,20 @@ async def inspect_model_request(
             )
         )
 
+    # convert openai tool choice to inspect tool_choice
+    inspect_tool_choice: ToolChoice | None = None
+    tool_choice: ChatCompletionToolChoiceOptionParam | None = json_data.get(
+        "tool_choice", None
+    )
+    if tool_choice is not None:
+        match tool_choice:
+            case "auto" | "none":
+                inspect_tool_choice = tool_choice
+            case "required":
+                inspect_tool_choice = "any"
+            case _:
+                inspect_tool_choice = ToolFunction(name=tool_choice["function"]["name"])
+
     # resolve model
     if model_name == "inspect":
         model = get_model()
@@ -122,6 +139,7 @@ async def inspect_model_request(
     output = await model.generate(
         input=input,
         tools=inspect_tools,
+        tool_choice=inspect_tool_choice,
         config=generate_config_from_openai(options),
     )
 
@@ -165,4 +183,16 @@ def generate_config_from_openai(options: FinalRequestOptions) -> GenerateConfig:
     config.parallel_tool_calls = json_data.get("parallel_tool_calls", None)
     config.reasoning_effort = json_data.get("reasoning_effort", None)
 
+    # response format
+    response_format: dict[str, Any] | None = json_data.get("response_format", None)
+    if response_format is not None:
+        json_schema: dict[str, Any] | None = response_format.get("json_schema", None)
+        if json_schema is not None:
+            config.response_schema = ResponseSchema(
+                name=json_schema.get("name", "schema"),
+                description=json_schema.get("description", None),
+                json_schema=JSONSchema.model_validate(json_schema.get("schema", {})),
+                strict=json_schema.get("strict", None),
+            )
+
     return config

inspect_ai/solver/_fork.py

@@ -1,10 +1,11 @@
-import asyncio
+import functools
 from contextvars import ContextVar
 from copy import deepcopy
 from typing import Any, cast
 
 from typing_extensions import overload
 
+from inspect_ai._util._async import tg_collect
 from inspect_ai._util.registry import registry_log_name, registry_params
 from inspect_ai.util._subtask import subtask
 
@@ -44,8 +45,9 @@ async def fork(
     if isinstance(solvers, Solver):
         return await solver_subtask(state, solvers)
     else:
-        subtasks = [solver_subtask(state, solver) for solver in solvers]
-        return await asyncio.gather(*subtasks)
+        return await tg_collect(
+            [functools.partial(solver_subtask, state, solver) for solver in solvers]
+        )
 
 
 async def solver_subtask(state: TaskState, solver: Solver) -> TaskState:

inspect_ai/solver/_human_agent/agent.py

@@ -1,6 +1,7 @@
-import asyncio
 from typing import cast
 
+import anyio
+
 from inspect_ai.util import display_type, input_panel, sandbox
 from inspect_ai.util._sandbox.events import SandboxEnvironmentProxy
 
@@ -42,7 +43,7 @@ def human_agent(
        Solver: Human agent solver.
     """
     # we can only run one human agent interaction at a time (use lock to enforce)
-    agent_lock = asyncio.Lock()
+    agent_lock = anyio.Lock()
 
     async def solve(state: TaskState, generate: Generate) -> TaskState:
         async with agent_lock:

inspect_ai/tool/__init__.py

@@ -20,7 +20,7 @@ from ._tool_call import (
 from ._tool_choice import ToolChoice, ToolFunction
 from ._tool_def import ToolDef
 from ._tool_info import ToolInfo
-from ._tool_params import JSONType, ToolParam, ToolParams
+from ._tool_params import ToolParam, ToolParams
 from ._tool_with import tool_with
 from ._tools._computer import computer
 from ._tools._execute import bash, python
@@ -56,12 +56,18 @@ __all__ = [
     "ToolInfo",
     "ToolParam",
     "ToolParams",
-    "JSONType",
 ]
 
 _UTIL_MODULE_VERSION = "0.3.19"
+_JSON_MODULE_VERSION = "0.3.73"
 _REMOVED_IN = "0.4"
 
+relocated_module_attribute(
+    "JSONType",
+    "inspect_ai.util.JSONType",
+    _JSON_MODULE_VERSION,
+    _REMOVED_IN,
+)
 
 relocated_module_attribute(
     "ToolEnvironment",

inspect_ai/tool/_tool_info.py

@@ -1,27 +1,19 @@
 import inspect
-import types
-import typing
-from dataclasses import is_dataclass
 from typing import (
     Any,
     Callable,
     Dict,
-    List,
-    Optional,
-    Tuple,
-    Type,
-    Union,
     get_args,
-    get_origin,
     get_type_hints,
-    is_typeddict,
 )
 
 from docstring_parser import Docstring, parse
 from pydantic import BaseModel, Field
 
+from inspect_ai.util._json import JSONType, json_schema
+
 from ._tool_description import tool_description
-from ._tool_params import JSONType, ToolParam, ToolParams
+from ._tool_params import ToolParam, ToolParams
 
 
 class ToolInfo(BaseModel):
@@ -88,7 +80,7 @@ def parse_tool_info(func: Callable[..., Any]) -> ToolInfo:
 
         # Get type information from type annotations
         if param_name in type_hints:
-            tool_param = parse_type(type_hints[param_name])
+            tool_param = json_schema(type_hints[param_name])
         # as a fallback try to parse it from the docstring
         # (this is minimally necessary for backwards compatiblity
         #  with tools gen1 type parsing, which only used docstrings)
@@ -129,84 +121,6 @@ def parse_tool_info(func: Callable[..., Any]) -> ToolInfo:
     return info
 
 
-def parse_type(type_hint: Type[Any]) -> ToolParam:
-    origin = get_origin(type_hint)
-    args = get_args(type_hint)
-
-    if origin is None:
-        if type_hint is int:
-            return ToolParam(type="integer")
-        elif type_hint is float:
-            return ToolParam(type="number")
-        elif type_hint is str:
-            return ToolParam(type="string")
-        elif type_hint is bool:
-            return ToolParam(type="boolean")
-        elif type_hint is list:
-            return ToolParam(type="array", items=ToolParam())
-        elif type_hint is dict:
-            return ToolParam(type="object", additionalProperties=ToolParam())
-        elif (
-            is_dataclass(type_hint)
-            or is_typeddict(type_hint)
-            or (isinstance(type_hint, type) and issubclass(type_hint, BaseModel))
-        ):
-            return parse_object(type_hint)
-        elif type_hint is type(None):
-            return ToolParam(type="null")
-        else:
-            return ToolParam()
-    elif origin is list or origin is List or origin is tuple or origin is Tuple:
-        return ToolParam(
-            type="array", items=parse_type(args[0]) if args else ToolParam()
-        )
-    elif origin is dict or origin is Dict:
-        return ToolParam(
-            type="object",
-            additionalProperties=parse_type(args[1]) if len(args) > 1 else ToolParam(),
-        )
-    elif origin is Union or origin is types.UnionType:
-        return ToolParam(anyOf=[parse_type(arg) for arg in args])
-    elif origin is Optional:
-        return ToolParam(
-            anyOf=[parse_type(arg) for arg in args] + [ToolParam(type="null")]
-        )
-    elif origin is typing.Literal:
-        return ToolParam(enum=list(args))
-
-    return ToolParam()  # Default case if we can't determine the type
-
-
-def parse_object(cls: Type[Any]) -> ToolParam:
-    properties: Dict[str, ToolParam] = {}
-    required: List[str] = []
-
-    if is_dataclass(cls):
-        fields = cls.__dataclass_fields__  # type: ignore
-        for name, field in fields.items():
-            properties[name] = parse_type(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] = ToolParam(**prop)
-        required = schema.get("required", [])
-    elif is_typeddict(cls):
-        annotations = get_type_hints(cls)
-        for name, type_hint in annotations.items():
-            properties[name] = parse_type(type_hint)
-            if name in cls.__required_keys__:
-                required.append(name)
-
-    return ToolParam(
-        type="object",
-        properties=properties,
-        required=required if required else None,
-        additionalProperties=False,
-    )
-
-
 def parse_docstring(docstring: str | None, param_name: str) -> Dict[str, str]:
     if not docstring:
         return {}