inspect-ai 0.3.69__py3-none-any.whl → 0.3.71__py3-none-any.whl

inspect_ai/model/_providers/hf.py

@@ -4,6 +4,7 @@ import functools
 import gc
 import json
 import os
+import time
 from dataclasses import dataclass
 from queue import Empty, Queue
 from threading import Thread
@@ -220,6 +221,7 @@ class HuggingFaceAPI(ModelAPI):
                 output_tokens=response.output_tokens,
                 total_tokens=response.total_tokens,
             ),
+            time=response.time,
         )
 
     @override
@@ -377,6 +379,7 @@ class GenerateOutput:
     output_tokens: int
     total_tokens: int
     logprobs: torch.Tensor | None
+    time: float
 
 
 @dataclass
@@ -432,6 +435,7 @@ def process_batches() -> None:
 
         try:
             # capture the generator and decoder functions
+            start_time = time.monotonic()
             first_input = inputs[0][0]
             device = first_input.device
             tokenizer = first_input.tokenizer
@@ -467,6 +471,7 @@ def process_batches() -> None:
             outputs = decoder(sequences=generated_tokens)
 
             # call back futures
+            total_time = time.monotonic() - start_time
             for i, output in enumerate(outputs):
                 future = inputs[i][1]
                 input_tokens = input_ids.size(dim=1)
@@ -483,6 +488,7 @@ def process_batches() -> None:
                         output_tokens=output_tokens,
                         total_tokens=input_tokens + output_tokens,
                         logprobs=logprobs[i] if logprobs is not None else None,
+                        time=total_time,
                     ),
                 )
 

inspect_ai/model/_providers/mistral.py

@@ -61,6 +61,7 @@ from .._model_output import (
     StopReason,
 )
 from .util import environment_prerequisite_error, model_base_url
+from .util.tracker import HttpxTimeTracker
 
 AZURE_MISTRAL_API_KEY = "AZURE_MISTRAL_API_KEY"
 AZUREAI_MISTRAL_API_KEY = "AZUREAI_MISTRAL_API_KEY"
@@ -125,57 +126,83 @@ class MistralAPI(ModelAPI):
         tool_choice: ToolChoice,
         config: GenerateConfig,
     ) -> ModelOutput | tuple[ModelOutput | Exception, ModelCall]:
-        # build request
-        request: dict[str, Any] = dict(
-            model=self.model_name,
-            messages=await mistral_chat_messages(input),
-            tools=mistral_chat_tools(tools) if len(tools) > 0 else None,
-            tool_choice=(
-                mistral_chat_tool_choice(tool_choice) if len(tools) > 0 else None
-            ),
-        )
-        if config.temperature is not None:
-            request["temperature"] = config.temperature
-        if config.top_p is not None:
-            request["top_p"] = config.top_p
-        if config.max_tokens is not None:
-            request["max_tokens"] = config.max_tokens
-        if config.seed is not None:
-            request["random_seed"] = config.seed
-
-        # send request
-        try:
-            with Mistral(
-                api_key=self.api_key,
-                timeout_ms=(config.timeout if config.timeout else DEFAULT_TIMEOUT)
-                * 1000,
-                **self.model_args,
-            ) as client:
-                response = await client.chat.complete_async(**request)
-        except SDKError as ex:
-            if ex.status_code == 400:
-                return self.handle_bad_request(ex), mistral_model_call(request, None)
-            else:
-                raise ex
-
-        if response is None:
-            raise RuntimeError("Mistral model did not return a response from generate.")
-
-        # return model output (w/ tool calls if they exist)
-        choices = completion_choices_from_response(response, tools)
-        return ModelOutput(
-            model=response.model,
-            choices=choices,
-            usage=ModelUsage(
-                input_tokens=response.usage.prompt_tokens,
-                output_tokens=(
-                    response.usage.completion_tokens
-                    if response.usage.completion_tokens
-                    else response.usage.total_tokens - response.usage.prompt_tokens
+        # create client
+        with Mistral(
+            api_key=self.api_key,
+            timeout_ms=(config.timeout if config.timeout else DEFAULT_TIMEOUT) * 1000,
+            **self.model_args,
+        ) as client:
+            # create time tracker
+            time_tracker = HttpxTimeTracker(client.sdk_configuration.async_client)
+
+            # build request
+            request_id = time_tracker.start_request()
+            request: dict[str, Any] = dict(
+                model=self.model_name,
+                messages=await mistral_chat_messages(input),
+                tools=mistral_chat_tools(tools) if len(tools) > 0 else None,
+                tool_choice=(
+                    mistral_chat_tool_choice(tool_choice) if len(tools) > 0 else None
                 ),
-                total_tokens=response.usage.total_tokens,
-            ),
-        ), mistral_model_call(request, response)
+                http_headers={HttpxTimeTracker.REQUEST_ID_HEADER: request_id},
+            )
+            if config.temperature is not None:
+                request["temperature"] = config.temperature
+            if config.top_p is not None:
+                request["top_p"] = config.top_p
+            if config.max_tokens is not None:
+                request["max_tokens"] = config.max_tokens
+            if config.seed is not None:
+                request["random_seed"] = config.seed
+
+            # prepare response for inclusion in model call
+            response: dict[str, Any] = {}
+
+            def model_call() -> ModelCall:
+                req = request.copy()
+                req.update(
+                    messages=[message.model_dump() for message in req["messages"]]
+                )
+                if req.get("tools", None) is not None:
+                    req["tools"] = [tool.model_dump() for tool in req["tools"]]
+
+                return ModelCall.create(
+                    request=req,
+                    response=response,
+                    time=time_tracker.end_request(request_id),
+                )
+
+            # send request
+            try:
+                completion = await client.chat.complete_async(**request)
+                response = completion.model_dump()
+            except SDKError as ex:
+                if ex.status_code == 400:
+                    return self.handle_bad_request(ex), model_call()
+                else:
+                    raise ex
+
+            if completion is None:
+                raise RuntimeError(
+                    "Mistral model did not return a response from generate."
+                )
+
+            # return model output (w/ tool calls if they exist)
+            choices = completion_choices_from_response(completion, tools)
+            return ModelOutput(
+                model=completion.model,
+                choices=choices,
+                usage=ModelUsage(
+                    input_tokens=completion.usage.prompt_tokens,
+                    output_tokens=(
+                        completion.usage.completion_tokens
+                        if completion.usage.completion_tokens
+                        else completion.usage.total_tokens
+                        - completion.usage.prompt_tokens
+                    ),
+                    total_tokens=completion.usage.total_tokens,
+                ),
+            ), model_call()
 
     @override
     def is_rate_limit(self, ex: BaseException) -> bool:
@@ -207,7 +234,7 @@ def mistral_model_call(
     request.update(messages=[message.model_dump() for message in request["messages"]])
     if request.get("tools", None) is not None:
         request["tools"] = [tool.model_dump() for tool in request["tools"]]
-    return ModelCall(
+    return ModelCall.create(
         request=request, response=response.model_dump() if response else {}
     )
 

inspect_ai/model/_providers/openai.py

@@ -1,8 +1,12 @@
 import os
+import socket
 from logging import getLogger
 from typing import Any
 
+import httpx
 from openai import (
+    DEFAULT_CONNECTION_LIMITS,
+    DEFAULT_TIMEOUT,
     APIConnectionError,
     APITimeoutError,
     AsyncAzureOpenAI,
@@ -21,6 +25,7 @@ from inspect_ai._util.constants import DEFAULT_MAX_RETRIES
 from inspect_ai._util.error import PrerequisiteError
 from inspect_ai._util.logger import warn_once
 from inspect_ai.model._openai import chat_choices_from_openai
+from inspect_ai.model._providers.util.tracker import HttpxTimeTracker
 from inspect_ai.tool import ToolChoice, ToolInfo
 
 from .._chat_message import ChatMessage
@@ -101,6 +106,9 @@ class OpenAIAPI(ModelAPI):
                         ],
                     )
 
+        # create async http client
+        http_client = OpenAIAsyncHttpxClient()
+
         # azure client
         if self.is_azure():
             # resolve base_url
@@ -125,6 +133,7 @@ class OpenAIAPI(ModelAPI):
                 max_retries=(
                     config.max_retries if config.max_retries else DEFAULT_MAX_RETRIES
                 ),
+                http_client=http_client,
                 **model_args,
             )
         else:
@@ -134,9 +143,13 @@ class OpenAIAPI(ModelAPI):
                 max_retries=(
                     config.max_retries if config.max_retries else DEFAULT_MAX_RETRIES
                 ),
+                http_client=http_client,
                 **model_args,
             )
 
+        # create time tracker
+        self._time_tracker = HttpxTimeTracker(self.client._client)
+
     def is_azure(self) -> bool:
         return self.service == "azure"
 
@@ -172,6 +185,9 @@ class OpenAIAPI(ModelAPI):
                 **self.completion_params(config, False),
             )
 
+        # allocate request_id (so we can see it from ModelCall)
+        request_id = self._time_tracker.start_request()
+
         # setup request and response for ModelCall
         request: dict[str, Any] = {}
         response: dict[str, Any] = {}
@@ -181,6 +197,7 @@ class OpenAIAPI(ModelAPI):
                 request=request,
                 response=response,
                 filter=image_url_filter,
+                time=self._time_tracker.end_request(request_id),
             )
 
         # unlike text models, vision models require a max_tokens (and set it to a very low
@@ -199,6 +216,7 @@ class OpenAIAPI(ModelAPI):
             tool_choice=openai_chat_tool_choice(tool_choice)
             if len(tools) > 0
             else NOT_GIVEN,
+            extra_headers={HttpxTimeTracker.REQUEST_ID_HEADER: request_id},
             **self.completion_params(config, len(tools) > 0),
         )
 
@@ -222,6 +240,16 @@ class OpenAIAPI(ModelAPI):
                     ModelUsage(
                         input_tokens=completion.usage.prompt_tokens,
                         output_tokens=completion.usage.completion_tokens,
+                        input_tokens_cache_read=(
+                            completion.usage.prompt_tokens_details.cached_tokens
+                            if completion.usage.prompt_tokens_details is not None
+                            else None  # openai only have cache read stats/pricing.
+                        ),
+                        reasoning_tokens=(
+                            completion.usage.completion_tokens_details.reasoning_tokens
+                            if completion.usage.completion_tokens_details is not None
+                            else None
+                        ),
                         total_tokens=completion.usage.total_tokens,
                     )
                     if completion.usage
@@ -241,10 +269,8 @@ class OpenAIAPI(ModelAPI):
     def is_rate_limit(self, ex: BaseException) -> bool:
         if isinstance(ex, RateLimitError):
             # Do not retry on these rate limit errors
-            if (
-                "Request too large" not in ex.message
-                and "You exceeded your current quota" not in ex.message
-            ):
+            # The quota exceeded one is related to monthly account quotas.
+            if "You exceeded your current quota" not in ex.message:
                 return True
         elif isinstance(
             ex, (APIConnectionError | APITimeoutError | InternalServerError)
@@ -333,3 +359,39 @@ class OpenAIAPI(ModelAPI):
             )
         else:
             return e
+
+
+class OpenAIAsyncHttpxClient(httpx.AsyncClient):
+    """Custom async client that deals better with long running Async requests.
+
+    Based on Anthropic DefaultAsyncHttpClient implementation that they
+    released along with Claude 3.7 as well as the OpenAI DefaultAsyncHttpxClient
+
+    """
+
+    def __init__(self, **kwargs: Any) -> None:
+        # This is based on the openai DefaultAsyncHttpxClient:
+        # https://github.com/openai/openai-python/commit/347363ed67a6a1611346427bb9ebe4becce53f7e
+        kwargs.setdefault("timeout", DEFAULT_TIMEOUT)
+        kwargs.setdefault("limits", DEFAULT_CONNECTION_LIMITS)
+        kwargs.setdefault("follow_redirects", True)
+
+        # This is based on the anthrpopic changes for claude 3.7:
+        # https://github.com/anthropics/anthropic-sdk-python/commit/c5387e69e799f14e44006ea4e54fdf32f2f74393#diff-3acba71f89118b06b03f2ba9f782c49ceed5bb9f68d62727d929f1841b61d12bR1387-R1403
+
+        # set socket options to deal with long running reasoning requests
+        socket_options = [
+            (socket.SOL_SOCKET, socket.SO_KEEPALIVE, True),
+            (socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 60),
+            (socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 5),
+        ]
+        TCP_KEEPIDLE = getattr(socket, "TCP_KEEPIDLE", None)
+        if TCP_KEEPIDLE is not None:
+            socket_options.append((socket.IPPROTO_TCP, TCP_KEEPIDLE, 60))
+
+        kwargs["transport"] = httpx.AsyncHTTPTransport(
+            limits=DEFAULT_CONNECTION_LIMITS,
+            socket_options=socket_options,
+        )
+
+        super().__init__(**kwargs)

inspect_ai/model/_providers/openai_o1.py

@@ -69,6 +69,16 @@ async def generate_o1(
         usage=ModelUsage(
             input_tokens=completion.usage.prompt_tokens,
             output_tokens=completion.usage.completion_tokens,
+            input_tokens_cache_read=(
+                completion.usage.prompt_tokens_details.cached_tokens
+                if completion.usage.prompt_tokens_details is not None
+                else None  # openai only have cache read stats/pricing.
+            ),
+            reasoning_tokens=(
+                completion.usage.completion_tokens_details.reasoning_tokens
+                if completion.usage.completion_tokens_details is not None
+                else None
+            ),
             total_tokens=completion.usage.total_tokens,
         )
         if completion.usage

inspect_ai/model/_providers/providers.py

@@ -48,7 +48,7 @@ def openai() -> type[ModelAPI]:
 def anthropic() -> type[ModelAPI]:
     FEATURE = "Anthropic API"
     PACKAGE = "anthropic"
-    MIN_VERSION = "0.29.0"
+    MIN_VERSION = "0.47.1"
 
     # verify we have the package
     try:
@@ -148,7 +148,7 @@ def cf() -> type[ModelAPI]:
 def mistral() -> type[ModelAPI]:
     FEATURE = "Mistral API"
     PACKAGE = "mistralai"
-    MIN_VERSION = "1.2.0"
+    MIN_VERSION = "1.5.0"
 
     # verify we have the package
     try:

inspect_ai/model/_providers/util/tracker.py

@@ -0,0 +1,92 @@
+import re
+import time
+from typing import Any, cast
+
+import httpx
+from shortuuid import uuid
+
+
+class HttpTimeTracker:
+    def __init__(self) -> None:
+        # track request start times
+        self._requests: dict[str, float] = {}
+
+    def start_request(self) -> str:
+        request_id = uuid()
+        self._requests[request_id] = time.monotonic()
+        return request_id
+
+    def end_request(self, request_id: str) -> float:
+        # read the request time if (if available) and purge from dict
+        request_time = self._requests.pop(request_id, None)
+        if request_time is None:
+            raise RuntimeError(f"request_id not registered: {request_id}")
+
+        # return elapsed time
+        return time.monotonic() - request_time
+
+    def update_request_time(self, request_id: str) -> None:
+        request_time = self._requests.get(request_id, None)
+        if not request_time:
+            raise RuntimeError(f"No request registered for request_id: {request_id}")
+
+        # update the request time
+        self._requests[request_id] = time.monotonic()
+
+
+class BotoTimeTracker(HttpTimeTracker):
+    def __init__(self, session: Any) -> None:
+        from aiobotocore.session import AioSession
+
+        super().__init__()
+
+        # register hook
+        session = cast(AioSession, session._session)
+        session.register(
+            "before-send.bedrock-runtime.Converse", self.converse_before_send
+        )
+
+    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 user_agent_extra(self, request_id: str) -> str:
+        return f"{self.USER_AGENT_PREFIX}{request_id}"
+
+    USER_AGENT_PREFIX = "ins/rid#"
+
+
+class HttpxTimeTracker(HttpTimeTracker):
+    """Class which tracks the duration of successful (200 status) http requests.
+
+    A special header is injected into requests which is then read from
+    an httpx '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.
+
+    To determine the total time, we also install an httpx response hook. In
+    this hook we look for 200 responses which have a registered request id.
+    When we find one, we update the end time of the request.
+
+    There is an 'end_request()' method which gets the total requeset time
+    for a request_id and then purges the request_id from our tracking (so
+    the dict doesn't grow unbounded)
+    """
+
+    REQUEST_ID_HEADER = "x-irid"
+
+    def __init__(self, client: httpx.AsyncClient):
+        super().__init__()
+
+        # install httpx request hook
+        client.event_hooks["request"].append(self.request_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)

inspect_ai/model/_providers/vllm.py

@@ -2,6 +2,7 @@ import asyncio
 import functools
 import gc
 import os
+import time
 from dataclasses import dataclass
 from queue import Empty, Queue
 from threading import Thread
@@ -48,7 +49,8 @@ class GenerateOutput:
     output_tokens: int
     total_tokens: int
     stop_reason: StopReason
-    logprobs: Logprobs | None = None
+    logprobs: Logprobs | None
+    time: float
 
 
 class VLLMAPI(ModelAPI):
@@ -258,6 +260,7 @@ class VLLMAPI(ModelAPI):
         ]
 
         # TODO: what's the best way to calculate token usage for num_choices > 1
+        total_time = responses[0].time
         input_tokens = responses[0].input_tokens
         output_tokens = sum(response.output_tokens for response in responses)
         total_tokens = input_tokens + output_tokens
@@ -270,6 +273,7 @@ class VLLMAPI(ModelAPI):
                 output_tokens=output_tokens,
                 total_tokens=total_tokens,
             ),
+            time=total_time,
         )
 
 
@@ -356,7 +360,7 @@ def get_stop_reason(finish_reason: str | None) -> StopReason:
 
 
 def post_process_output(
-    output: RequestOutput, i: int, num_top_logprobs: int | None
+    output: RequestOutput, i: int, num_top_logprobs: int | None, total_time: float
 ) -> GenerateOutput:
     completion = output.outputs[i]
     output_text: str = completion.text
@@ -377,14 +381,15 @@ def post_process_output(
         total_tokens=total_tokens,
         stop_reason=get_stop_reason(completion.finish_reason),
         logprobs=extract_logprobs(completion, num_top_logprobs),
+        time=total_time,
     )
 
 
 def post_process_outputs(
-    output: RequestOutput, num_top_logprobs: int | None
+    output: RequestOutput, num_top_logprobs: int | None, total_time: float
 ) -> list[GenerateOutput]:
     return [
-        post_process_output(output, i, num_top_logprobs)
+        post_process_output(output, i, num_top_logprobs, total_time)
         for i in range(len(output.outputs))
     ]
 
@@ -412,6 +417,7 @@ def process_batches() -> None:
             continue
 
         try:
+            start_time = time.monotonic()
             first_input = inputs[0][0]
             generator = first_input.generator
             num_top_logprobs = first_input.num_top_logprobs
@@ -419,6 +425,7 @@ def process_batches() -> None:
             # generate
             outputs = generator([input[0].input for input in inputs])
 
+            total_time = time.monotonic() - start_time
             for i, output in enumerate(outputs):
                 future = inputs[i][1]
 
@@ -426,7 +433,8 @@ def process_batches() -> None:
                 # 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, post_process_outputs(output, num_top_logprobs)
+                    future.set_result,
+                    post_process_outputs(output, num_top_logprobs, total_time),
                 )
 
         except Exception as e:

inspect_ai/model/_reasoning.py

@@ -5,13 +5,26 @@ from typing import NamedTuple
 class ContentWithReasoning(NamedTuple):
     content: str
     reasoning: str
+    signature: str | None = None
+    redacted: bool = False
 
 
 def parse_content_with_reasoning(content: str) -> ContentWithReasoning | None:
-    match = re.match(r"\s*<think>(.*?)</think>(.*)", content, re.DOTALL)
+    # Match <think> tag with optional attributes
+    pattern = r'\s*<think(?:\s+signature="([^"]*)")?(?:\s+redacted="(true)")?\s*>(.*?)</think>(.*)'
+    match = re.match(pattern, content, re.DOTALL)
+
     if match:
+        signature = match.group(1)  # This will be None if not present
+        redacted_value = match.group(2)  # This will be "true" or None
+        reasoning = match.group(3).strip()
+        content_text = match.group(4).strip()
+
         return ContentWithReasoning(
-            content=match.group(2).strip(), reasoning=match.group(1).strip()
+            content=content_text,
+            reasoning=reasoning,
+            signature=signature,
+            redacted=redacted_value == "true",
         )
     else:
         return None

inspect_ai/scorer/_model.py

@@ -274,25 +274,29 @@ def chat_history(state: TaskState) -> str:
 
     # begin history with text of first message (it will come right after
     # 'Task' or 'Question' in the template)
-    history: list[str] = [messages[0].text]
-
-    # for subsequent messages present with e.g. Assistant: {message.text}
-    for message in messages[1:]:
-        if isinstance(message, ChatMessageUser):
-            history.append(f"User: {message.text}")
-        elif isinstance(message, ChatMessageAssistant):
-            assistant_message = [message.text] if message.text else []
-            if message.tool_calls:
-                assistant_message.extend(
-                    [
-                        format_function_call(tool_call.function, tool_call.arguments)
-                        for tool_call in message.tool_calls
-                    ]
+    history: list[str] = []
+    if len(messages) > 0:
+        history.append(messages[0].text)
+
+        # for subsequent messages present with e.g. Assistant: {message.text}
+        for message in messages[1:]:
+            if isinstance(message, ChatMessageUser):
+                history.append(f"User: {message.text}")
+            elif isinstance(message, ChatMessageAssistant):
+                assistant_message = [message.text] if message.text else []
+                if message.tool_calls:
+                    assistant_message.extend(
+                        [
+                            format_function_call(
+                                tool_call.function, tool_call.arguments
+                            )
+                            for tool_call in message.tool_calls
+                        ]
+                    )
+                history.append("Assistant: " + "\n\n".join(assistant_message))
+            elif isinstance(message, ChatMessageTool):
+                history.append(
+                    f"Tool ({message.function}): {message.tool_error or ''}{message.text}"
                 )
-            history.append("Assistant: " + "\n\n".join(assistant_message))
-        elif isinstance(message, ChatMessageTool):
-            history.append(
-                f"Tool ({message.function}): {message.tool_error or ''}{message.text}"
-            )
 
     return "\n\n".join(history)

inspect_ai/solver/_basic_agent.py

@@ -24,7 +24,7 @@ logger = getLogger(__name__)
 
 DEFAULT_SYSTEM_MESSAGE = """
 You are a helpful assistant attempting to submit the correct answer. You have
-several functions available to help with finding the answer. Each message may
+several functions available to help with finding the answer. Each message
 may perform one function call. You will see the result of the function right
 after sending the message. If you need to perform multiple actions, you can
 always send more messages with subsequent function calls. Do some reasoning
@@ -206,13 +206,11 @@ def basic_agent(
                             # exit if we are at max_attempts
                             attempts += 1
                             if attempts >= max_attempts:
-                                state.completed = True
                                 break
 
                             # exit if the submission is successful
                             answer_scores = await score(state)
                             if score_value_fn(answer_scores[0].value) == 1.0:
-                                state.completed = True
                                 break
 
                             # otherwise notify the model that it was incorrect and continue

inspect_ai/solver/_bridge/patch.py

@@ -72,8 +72,6 @@ def init_openai_request_patch() -> None:
                 _patch_enabled.get()
                 # completions request
                 and options.url == "/chat/completions"
-                # call to openai not another service (e.g. TogetherAI)
-                and self.base_url == "https://api.openai.com/v1/"
             ):
                 # must also be an explicit request for an inspect model
                 json_data = cast(dict[str, Any], options.json_data)