inspect-ai 0.3.99__py3-none-any.whl → 0.3.100__py3-none-any.whl

inspect_ai/model/_openai_responses.py

@@ -1,6 +1,5 @@
 import json
-from itertools import chain
-from typing import TypedDict, cast
+from typing import Sequence, TypedDict, cast
 
 from openai.types.responses import (
     FunctionToolParam,
@@ -8,6 +7,8 @@ from openai.types.responses import (
     ResponseComputerToolCallParam,
     ResponseFunctionToolCall,
     ResponseFunctionToolCallParam,
+    ResponseFunctionWebSearch,
+    ResponseFunctionWebSearchParam,
     ResponseInputContentParam,
     ResponseInputImageParam,
     ResponseInputItemParam,
@@ -51,6 +52,7 @@ from inspect_ai.model._openai_computer_use import (
     maybe_computer_use_preview_tool,
     tool_call_from_openai_computer_tool_call,
 )
+from inspect_ai.model._openai_web_search import maybe_web_search_tool
 from inspect_ai.tool._tool_call import ToolCall
 from inspect_ai.tool._tool_choice import ToolChoice
 from inspect_ai.tool._tool_info import ToolInfo
@@ -174,6 +176,12 @@ def openai_responses_chat_choices(
     return [ChatCompletionChoice(message=message, stop_reason=stop_reason)]
 
 
+def is_native_tool_configured(
+    tools: Sequence[ToolInfo], config: GenerateConfig
+) -> bool:
+    return any(_maybe_native_tool_param(tool, config) is not None for tool in tools)
+
+
 # The next two function perform transformations between OpenAI types an Inspect
 # ChatMessageAssistant. Here is a diagram that helps visualize the transforms.
 # ┌───────────────────────────┐    ┌───────────────────────────┐    ┌───────────────────────────┐
@@ -207,7 +215,6 @@ def openai_responses_chat_choices(
 
 
 class _AssistantInternal(TypedDict):
-    output_message_id: str | None
     tool_message_ids: dict[str, str]
 
 
@@ -237,17 +244,17 @@ def _chat_message_assistant_from_openai_response(
     # collect output and tool calls
     message_content: list[Content] = []
     tool_calls: list[ToolCall] = []
-    internal = _AssistantInternal(output_message_id=None, tool_message_ids={})
+    internal = _AssistantInternal(tool_message_ids={})
     for output in response.output:
         match output:
             case ResponseOutputMessage(content=content, id=id):
-                assert internal["output_message_id"] is None, "Multiple message outputs"
-                internal["output_message_id"] = id
                 message_content.extend(
                     [
-                        ContentText(text=c.text)
+                        ContentText(text=c.text, internal={"id": id})
                         if isinstance(c, ResponseOutputText)
-                        else ContentText(text=c.refusal, refusal=True)
+                        else ContentText(
+                            text=c.refusal, refusal=True, internal={"id": id}
+                        )
                         for c in content
                     ]
                 )
@@ -277,6 +284,13 @@ def _chat_message_assistant_from_openai_response(
                         tool_calls.append(
                             tool_call_from_openai_computer_tool_call(output)
                         )
+                    case ResponseFunctionWebSearch():
+                        # We don't currently capture this since the model did the
+                        # "tool call" internally. It's conceivable that could be
+                        # forced to include it in `.internal` in the future, but
+                        # for now we just ignore it.
+                        # {"id":"ws_682cdcec3fa88198bc10b38fafefbd5e077e89e31fd4a3d5","status":"completed","type":"web_search_call"}
+                        pass
                     case _:
                         raise ValueError(f"Unexpected output type: {output.__class__}")
 
@@ -304,25 +318,39 @@ def _openai_input_items_from_chat_message_assistant(
     field of the `ChatMessageAssistant` to help it provide the proper id's the
     items in the returned list.
     """
-    (output_message_id, tool_message_ids) = _ids_from_assistant_internal(message)
+    tool_message_ids = _ids_from_assistant_internal(message)
 
     # we want to prevent yielding output messages in the case where we have an
     # 'internal' field (so the message came from the model API as opposed to
-    # being user synthesized) AND there is no output_message_id (indicating that
-    # when reading the message from the server we didn't find output). this could
-    # happen e.g. when a react() agent sets the output.completion in response
+    # being user synthesized) AND there are no ContentText items with message IDs
+    # (indicating that when reading the message from the server we didn't find output).
+    # this could happen e.g. when a react() agent sets the output.completion in response
     # to a submit() tool call
-    suppress_output_message = message.internal is not None and output_message_id is None
+    content_items: list[ContentText | ContentReasoning] = (
+        [ContentText(text=message.content)]
+        if isinstance(message.content, str)
+        else [
+            c for c in message.content if isinstance(c, ContentText | ContentReasoning)
+        ]
+    )
+    has_content_with_ids = any(
+        isinstance(c, ContentText)
+        and isinstance(c.internal, dict)
+        and "id" in c.internal
+        for c in content_items
+    )
+    suppress_output_message = message.internal is not None and not has_content_with_ids
 
     # if we are not storing messages on the server then blank these out
     if not store:
-        output_message_id = None
         tool_message_ids = {}
 
-    # items to return -- ensure we use a single output message (and just chain
-    # additional content on to it)
+    # items to return
     items: list[ResponseInputItemParam] = []
-    output_message: ResponseOutputMessageParam | None = None
+    # group content by message ID
+    messages_by_id: dict[
+        str | None, list[ResponseOutputTextParam | ResponseOutputRefusalParam]
+    ] = {}
 
     for content in (
         list[ContentText | ContentReasoning]([ContentText(text=message.content)])
@@ -352,6 +380,14 @@ def _openai_input_items_from_chat_message_assistant(
                 if suppress_output_message:
                     continue
 
+                # get the message ID from ContentText.modelJson
+                content_message_id: str | None = None
+                if isinstance(content.internal, dict) and "id" in content.internal:
+                    id_value = content.internal["id"]
+                    content_message_id = id_value if isinstance(id_value, str) else None
+                else:
+                    content_message_id = None
+
                 new_content = (
                     ResponseOutputRefusalParam(type="refusal", refusal=text)
                     if refusal
@@ -359,22 +395,24 @@ def _openai_input_items_from_chat_message_assistant(
                         type="output_text", text=text, annotations=[]
                     )
                 )
-                if output_message is None:
-                    output_message = ResponseOutputMessageParam(
-                        type="message",
-                        role="assistant",
-                        # this actually can be `None`, and it will in fact be `None` when the
-                        # assistant message is synthesized by the scaffold as opposed to being
-                        # replayed from the model (or when store=False)
-                        id=output_message_id,  # type: ignore[typeddict-item]
-                        content=[new_content],
-                        status="completed",
-                    )
-                    items.append(output_message)
-                else:
-                    output_message["content"] = chain(
-                        output_message["content"], [new_content]
-                    )
+
+                if content_message_id not in messages_by_id:
+                    messages_by_id[content_message_id] = []
+                messages_by_id[content_message_id].append(new_content)
+
+    # create ResponseOutputMessage for each unique ID
+    for msg_id, content_list in messages_by_id.items():
+        output_message = ResponseOutputMessageParam(
+            type="message",
+            role="assistant",
+            # this actually can be `None`, and it will in fact be `None` when the
+            # assistant message is synthesized by the scaffold as opposed to being
+            # replayed from the model (or when store=False)
+            id=msg_id,  # type: ignore[typeddict-item]
+            content=content_list,
+            status="completed",
+        )
+        items.append(output_message)
 
     return items + _tool_call_items_from_assistant_message(message, tool_message_ids)
 
@@ -399,7 +437,7 @@ def _maybe_native_tool_param(
 ) -> ToolParam | None:
     return (
         (
-            maybe_computer_use_preview_tool(tool)
+            maybe_computer_use_preview_tool(tool) or maybe_web_search_tool(tool)
             # or self.text_editor_tool_param(tool)
             # or self.bash_tool_param(tool)
         )
@@ -442,22 +480,23 @@ def _tool_call_items_from_assistant_message(
 
 def _ids_from_assistant_internal(
     message: ChatMessageAssistant,
-) -> tuple[str | None, dict[str, str]]:
+) -> dict[str, str]:
     if message.internal is not None:
         assert isinstance(message.internal, dict), (
             "OpenAI ChatMessageAssistant internal must be an _AssistantInternal"
         )
         internal = cast(_AssistantInternal, message.internal)
-        return (internal["output_message_id"], internal["tool_message_ids"])
+        return internal["tool_message_ids"]
     else:
-        return None, {}
+        return {}
 
 
 _ResponseToolCallParam = (
-    ResponseFunctionToolCallParam | ResponseComputerToolCallParam
+    ResponseFunctionToolCallParam
+    | ResponseComputerToolCallParam
+    | ResponseFunctionWebSearchParam
     # | ResponseFileSearchToolCallParam
     # | ResponseFunctionToolCallParam
-    # | ResponseFunctionWebSearchParam
 )
 
 

inspect_ai/model/_openai_web_search.py

@@ -0,0 +1,31 @@
+from typing import cast
+
+from openai.types.responses import WebSearchTool, WebSearchToolParam
+
+from inspect_ai.tool._tool_info import ToolInfo
+
+
+def maybe_web_search_tool(tool: ToolInfo) -> WebSearchToolParam | None:
+    return (
+        _web_search_tool(tool.options["openai"])
+        if tool.name == "web_search" and tool.options and "openai" in tool.options
+        else None
+    )
+
+
+def _web_search_tool(maybe_openai_options: object) -> WebSearchToolParam:
+    if maybe_openai_options is None:
+        maybe_openai_options = {}
+    elif not isinstance(maybe_openai_options, dict):
+        raise TypeError(
+            f"Expected a dictionary for openai_options, got {type(maybe_openai_options)}"
+        )
+    openai_options = (
+        WebSearchTool.model_validate(
+            {"type": "web_search_preview", **maybe_openai_options}
+        )
+        if maybe_openai_options
+        else WebSearchTool(type="web_search_preview")
+    )
+
+    return cast(WebSearchToolParam, openai_options.model_dump(exclude_none=True))

inspect_ai/model/_providers/azureai.py

@@ -1,3 +1,4 @@
+import functools
 import json
 import os
 from copy import copy
@@ -151,7 +152,7 @@ class AzureAIAPI(ModelAPI):
 
         # prepare request
         request = dict(
-            messages=await chat_request_messages(input, handler),
+            messages=await chat_request_messages(input, handler, self.is_mistral()),
             **self.completion_params(config),
         )
         # newer versions of vllm reject requests with tools or tool_choice if the
@@ -280,9 +281,77 @@ class AzureAIAPI(ModelAPI):
 
 
 async def chat_request_messages(
-    messages: list[ChatMessage], handler: ChatAPIHandler | None
+    messages: list[ChatMessage],
+    handler: ChatAPIHandler | None,
+    is_mistral: bool = False,
+) -> list[ChatRequestMessage]:
+    chat_messages = [
+        await chat_request_message(message, handler) for message in messages
+    ]
+    if is_mistral:
+        chat_messages = functools.reduce(mistral_message_reducer, chat_messages, [])
+
+    return chat_messages
+
+
+def mistral_message_reducer(
+    messages: list[ChatRequestMessage],
+    message: ChatRequestMessage,
 ) -> list[ChatRequestMessage]:
-    return [await chat_request_message(message, handler) for message in messages]
+    """Fold any user messages found immediately after tool messages into the last tool message."""
+    if (
+        len(messages) > 0
+        and isinstance(messages[-1], ToolMessage)
+        and isinstance(message, UserMessage)
+    ):
+        messages[-1] = fold_user_message_into_tool_message(messages[-1], message)
+    else:
+        messages.append(message)
+
+    return messages
+
+
+def fold_user_message_into_tool_message(
+    tool_message: ToolMessage,
+    user_message: UserMessage,
+) -> ToolMessage:
+    def convert_content_items_to_string(list_content: list[ContentItem]) -> str:
+        if not all(
+            isinstance(item, (TextContentItem | ImageContentItem))
+            for item in list_content
+        ):
+            raise TypeError(
+                "Expected all items to be TextContentItem or ImageContentItem"
+            )
+
+        parts = []
+        for item in list_content:
+            if isinstance(item, TextContentItem):
+                parts.append(item.text)
+            elif isinstance(item, ImageContentItem):
+                parts.append(f"[Image: {item.image_url.url}]")
+            else:
+                raise ValueError("Unexpected content item type")
+        return "".join(parts)
+
+    def normalise_content(
+        content: str | list[ContentItem] | None,
+    ) -> str | None:
+        return (
+            None
+            if content is None
+            else convert_content_items_to_string(content)
+            if isinstance(content, list)
+            else content
+        )
+
+    tool_content = normalise_content(tool_message.content)
+    user_content = normalise_content(user_message.content)
+
+    return ToolMessage(
+        content=(tool_content or "") + (user_content or ""),
+        tool_call_id=tool_message.tool_call_id,
+    )
 
 
 async def chat_request_message(

inspect_ai/model/_providers/openai.py

@@ -42,6 +42,7 @@ from .._openai import (
     openai_media_filter,
     openai_should_retry,
 )
+from .._openai_responses import is_native_tool_configured
 from .openai_o1 import generate_o1
 from .util import environment_prerequisite_error, model_base_url
 
@@ -241,7 +242,7 @@ class OpenAIAPI(ModelAPI):
                 tools=tools,
                 **self.completion_params(config, False),
             )
-        elif self.responses_api:
+        elif self.responses_api or is_native_tool_configured(tools, config):
             return await generate_responses(
                 client=self.client,
                 http_hooks=self._http_hooks,

inspect_ai/scorer/_metric.py

@@ -7,7 +7,6 @@ from typing import (
     Protocol,
     Type,
     Union,
-    cast,
     overload,
     runtime_checkable,
 )
@@ -356,7 +355,7 @@ def metric(
             )
             return metric
 
-        return metric_register(cast(Callable[P, Metric], metric_wrapper), metric_name)
+        return metric_register(metric_wrapper, metric_name)
 
     # for decorators with an explicit name, one more wrapper for the name
     if isinstance(name, str):

inspect_ai/solver/_task_state.py

@@ -290,7 +290,7 @@ class TaskState:
         return self._tools
 
     @tools.setter
-    def tools(self, tools: list[Tool | ToolDef]) -> None:
+    def tools(self, tools: Sequence[Tool | ToolDef]) -> None:
         self._tools.clear()
         for tool in tools:
             self._tools.append(tool if isinstance(tool, Tool) else tool.as_tool())
@@ -353,7 +353,7 @@ class TaskState:
     def completed(self) -> bool:
         """Is the task completed.
 
-        Additionally, checks message and token limits and raises if they are exceeded, and also checks for an operator interrupt of the sample.
+        Additionally, checks for an operator interrupt of the sample.
         """
         from inspect_ai.log._samples import set_active_sample_total_messages
 

inspect_ai/tool/_tool.py

@@ -224,13 +224,15 @@ def tool(
             tool_parallel = parallel
             tool_viewer = viewer
             tool_model_input = model_input
+            tool_options: dict[str, object] | None = None
             if is_registry_object(tool):
-                _, _, reg_parallel, reg_viewer, reg_model_input = tool_registry_info(
-                    tool
+                _, _, reg_parallel, reg_viewer, reg_model_input, options = (
+                    tool_registry_info(tool)
                 )
                 tool_parallel = parallel and reg_parallel
                 tool_viewer = viewer or reg_viewer
                 tool_model_input = model_input or reg_model_input
+                tool_options = options
 
             # tag the object
             registry_tag(
@@ -247,6 +249,7 @@ def tool(
                             tool_model_input
                             or getattr(tool, TOOL_INIT_MODEL_INPUT, None)
                         ),
+                        TOOL_OPTIONS: tool_options,
                     },
                 ),
                 *args,
@@ -267,6 +270,7 @@ TOOL_PROMPT = "prompt"
 TOOL_PARALLEL = "parallel"
 TOOL_VIEWER = "viewer"
 TOOL_MODEL_INPUT = "model_input"
+TOOL_OPTIONS = "options"
 
 
 TOOL_INIT_MODEL_INPUT = "__TOOL_INIT_MODEL_INPUT__"

inspect_ai/tool/_tool_def.py

@@ -16,6 +16,7 @@ from inspect_ai._util.registry import (
 
 from ._tool import (
     TOOL_MODEL_INPUT,
+    TOOL_OPTIONS,
     TOOL_PARALLEL,
     TOOL_PROMPT,
     TOOL_VIEWER,
@@ -44,6 +45,7 @@ class ToolDef:
         parallel: bool | None = None,
         viewer: ToolCallViewer | None = None,
         model_input: ToolCallModelInput | None = None,
+        options: dict[str, object] | None = None,
     ) -> None:
         """Create a tool definition.
 
@@ -59,6 +61,8 @@ class ToolDef:
           viewer: Optional tool call viewer implementation.
           model_input: Optional function that determines how
               tool call results are played back as model input.
+          options: Optional property bag that can be used by the model provider
+              to customize the implementation of the tool
 
         Returns:
           Tool definition.
@@ -82,6 +86,7 @@ class ToolDef:
             self.parallel = parallel if parallel is not None else tdef.parallel
             self.viewer = viewer or tdef.viewer
             self.model_input = model_input or tdef.model_input
+            self.options = options or tdef.options
 
         # if its not a tool then extract tool_info if all fields have not
         # been provided explicitly
@@ -112,6 +117,7 @@ class ToolDef:
             self.parallel = parallel is not False
             self.viewer = viewer
             self.model_input = model_input
+            self.options = options
 
     tool: Callable[..., Any]
     """Callable to execute tool."""
@@ -134,13 +140,20 @@ class ToolDef:
     model_input: ToolCallModelInput | None
     """Custom model input presenter for tool calls."""
 
+    options: dict[str, object] | None = None
+    """Optional property bag that can be used by the model provider to customize the implementation of the tool"""
+
     def as_tool(self) -> Tool:
         """Convert a ToolDef to a Tool."""
         tool = self.tool
         info = RegistryInfo(
             type="tool",
             name=self.name,
-            metadata={TOOL_PARALLEL: self.parallel, TOOL_VIEWER: self.viewer},
+            metadata={
+                TOOL_PARALLEL: self.parallel,
+                TOOL_VIEWER: self.viewer,
+                TOOL_OPTIONS: self.options,
+            },
         )
         set_registry_info(tool, info)
         set_registry_params(tool, {})
@@ -189,11 +202,12 @@ class ToolDefFields(NamedTuple):
     parallel: bool
     viewer: ToolCallViewer | None
     model_input: ToolCallModelInput | None
+    options: dict[str, object] | None
 
 
 def tool_def_fields(tool: Tool) -> ToolDefFields:
     # get tool_info
-    name, prompt, parallel, viewer, model_input = tool_registry_info(tool)
+    name, prompt, parallel, viewer, model_input, options = tool_registry_info(tool)
     tool_info = parse_tool_info(tool)
 
     # if there is a description then append any prompt to the
@@ -234,19 +248,28 @@ def tool_def_fields(tool: Tool) -> ToolDefFields:
         parallel=parallel,
         viewer=viewer,
         model_input=model_input,
+        options=options,
     )
 
 
 def tool_registry_info(
     tool: Tool,
-) -> tuple[str, str | None, bool, ToolCallViewer | None, ToolCallModelInput | None]:
+) -> tuple[
+    str,
+    str | None,
+    bool,
+    ToolCallViewer | None,
+    ToolCallModelInput | None,
+    dict[str, object] | None,
+]:
     info = registry_info(tool)
     name = info.name.split("/")[-1]
     prompt = info.metadata.get(TOOL_PROMPT, None)
     parallel = info.metadata.get(TOOL_PARALLEL, True)
     viewer = info.metadata.get(TOOL_VIEWER, None)
     model_input = info.metadata.get(TOOL_MODEL_INPUT, None)
-    return name, prompt, parallel, viewer, model_input
+    options = info.metadata.get(TOOL_OPTIONS, None)
+    return name, prompt, parallel, viewer, model_input, options
 
 
 def validate_tool_parameters(tool_name: str, parameters: dict[str, ToolParam]) -> None:

inspect_ai/tool/_tool_info.py

@@ -49,6 +49,8 @@ class ToolInfo(BaseModel):
     """Short description of tool."""
     parameters: ToolParams = Field(default_factory=ToolParams)
     """JSON Schema of tool parameters object."""
+    options: dict[str, object] | None = Field(default=None)
+    """Optional property bag that can be used by the model provider to customize the implementation of the tool"""
 
 
 def parse_tool_info(func: Callable[..., Any]) -> ToolInfo:

inspect_ai/tool/_tools/_web_search/_google.py

@@ -4,6 +4,7 @@ from typing import Awaitable, Callable
 import anyio
 import httpx
 from bs4 import BeautifulSoup, NavigableString
+from pydantic import BaseModel
 from tenacity import (
     retry,
     retry_if_exception,
@@ -23,6 +24,13 @@ Page Content: {text}
 """
 
 
+class GoogleOptions(BaseModel):
+    num_results: int | None = None
+    max_provider_calls: int | None = None
+    max_connections: int | None = None
+    model: str | None = None
+
+
 class SearchLink:
     def __init__(self, url: str, snippet: str) -> None:
         self.url = url
@@ -42,11 +50,14 @@ def maybe_get_google_api_keys() -> tuple[str, str] | None:
 
 
 def google_search_provider(
-    num_results: int,
-    max_provider_calls: int,
-    max_connections: int,
-    model: str | None,
+    in_options: dict[str, object] | None = None,
 ) -> Callable[[str], Awaitable[str | None]]:
+    options = GoogleOptions.model_validate(in_options) if in_options else None
+    num_results = (options.num_results if options else None) or 3
+    max_provider_calls = (options.max_provider_calls if options else None) or 3
+    max_connections = (options.max_connections if options else None) or 10
+    model = options.model if options else None
+
     keys = maybe_get_google_api_keys()
     if not keys:
         raise PrerequisiteError(

inspect_ai/tool/_tools/_web_search/_tavily.py

@@ -1,5 +1,5 @@
 import os
-from typing import Awaitable, Callable
+from typing import Awaitable, Callable, Literal
 
 import httpx
 from pydantic import BaseModel, Field
@@ -16,6 +16,25 @@ from inspect_ai._util.httpx import httpx_should_retry, log_httpx_retry_attempt
 from inspect_ai.util._concurrency import concurrency
 
 
+class TavilyOptions(BaseModel):
+    topic: Literal["general", "news"] | None = None
+    search_depth: Literal["basic", "advanced"] | None = None
+    chunks_per_source: Literal[1, 2, 3] | None = None
+    max_results: int | None = None
+    time_range: Literal["day", "week", "month", "year", "d", "w", "m", "y"] | None = (
+        None
+    )
+    days: int | None = None
+    include_answer: bool | Literal["basic", "advanced"] | None = None
+    include_raw_content: bool | None = None
+    include_images: bool | None = None
+    include_image_descriptions: bool | None = None
+    include_domains: list[str] | None = None
+    exclude_domains: list[str] | None = None
+    # max_connections is not a Tavily API option, but an inspect option
+    max_connections: int | None = None
+
+
 class TavilySearchResult(BaseModel):
     title: str
     url: str
@@ -32,17 +51,25 @@ class TavilySearchResponse(BaseModel):
 
 
 def tavily_search_provider(
-    num_results: int, max_connections: int
+    in_options: dict[str, object] | None = None,
 ) -> Callable[[str], Awaitable[str | None]]:
+    options = TavilyOptions.model_validate(in_options) if in_options else None
+    # Separate max_connections (which is an inspect thing) from the rest of the
+    # options which will be passed in the request body
+    max_connections = (options.max_connections if options else None) or 10
+    api_options = (
+        options.model_dump(exclude={"max_connections"}, exclude_none=True)
+        if options
+        else {}
+    )
+    if not api_options.get("include_answer", False):
+        api_options["include_answer"] = True
+
     tavily_api_key = os.environ.get("TAVILY_API_KEY", None)
     if not tavily_api_key:
         raise PrerequisiteError(
             "TAVILY_API_KEY not set in the environment. Please ensure ths variable is defined to use Tavily with the web_search tool.\n\nLearn more about the Tavily web search provider at https://inspect.aisi.org.uk/tools.html#tavily-provider"
         )
-    if num_results > 20:
-        raise PrerequisiteError(
-            "The Tavily search provider is limited to 20 results per query."
-        )
 
     # Create the client within the provider
     client = httpx.AsyncClient(timeout=30)
@@ -52,12 +79,8 @@ def tavily_search_provider(
         headers = {
             "Authorization": f"Bearer {tavily_api_key}",
         }
-        body = {
-            "query": query,
-            "max_results": 10,  # num_results,
-            # "search_depth": "advanced",
-            "include_answer": "advanced",
-        }
+
+        body = {"query": query, **api_options}
 
         # retry up to 5 times over a period of up to 1 minute
         @retry(