inspect-ai 0.3.57__py3-none-any.whl → 0.3.59__py3-none-any.whl

inspect_ai/log/_transcript.py

@@ -1,3 +1,4 @@
+import asyncio
 import contextlib
 from contextvars import ContextVar
 from datetime import datetime
@@ -11,7 +12,7 @@ from typing import (
     Union,
 )
 
-from pydantic import BaseModel, Field, JsonValue, field_serializer
+from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_serializer
 
 from inspect_ai._util.constants import SAMPLE_SUBTASK
 from inspect_ai._util.error import EvalError
@@ -69,7 +70,7 @@ class SampleLimitEvent(BaseEvent):
     event: Literal["sample_limit"] = Field(default="sample_limit")
     """Event type."""
 
-    type: Literal["message", "time", "token", "operator"]
+    type: Literal["message", "time", "token", "operator", "custom"]
     """Type of limit that halted processing"""
 
     message: str
@@ -123,6 +124,9 @@ class ModelEvent(BaseEvent):
     output: ModelOutput
     """Output from model."""
 
+    error: str | None = Field(default=None)
+    """Error which occurred during model call."""
+
     cache: Literal["read", "write"] | None = Field(default=None)
     """Was this a cache read or write."""
 
@@ -176,6 +180,32 @@ class ToolEvent(BaseEvent):
         self.events = events
         self.pending = None
 
+    # mechanism for operator to cancel the tool call
+
+    def set_task(self, task: asyncio.Task[Any]) -> None:
+        """Set the tool task (for possible cancellation)"""
+        self._task = task
+
+    def cancel(self) -> None:
+        """Cancel the tool task."""
+        if self._task:
+            self._cancelled = True
+            self._task.cancel()
+
+    @property
+    def cancelled(self) -> bool:
+        """Was the task cancelled?"""
+        return self._cancelled is True
+
+    _cancelled: bool | None = None
+    """Was this tool call cancelled?"""
+
+    _task: asyncio.Task[Any] | None = None
+    """Handle to task (used for cancellation)"""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    """Required so that we can include '_task' as a member."""
+
 
 class ApprovalEvent(BaseEvent):
     """Tool approval."""

inspect_ai/model/__init__.py

@@ -1,6 +1,12 @@
 # ruff: noqa: F401 F403 F405
 
-from inspect_ai._util.content import Content, ContentImage, ContentText
+from inspect_ai._util.content import (
+    Content,
+    ContentAudio,
+    ContentImage,
+    ContentText,
+    ContentVideo,
+)
 from inspect_ai._util.deprecation import relocated_module_attribute
 
 from ._cache import (
@@ -42,8 +48,10 @@ __all__ = [
     "GenerateConfig",
     "GenerateConfigArgs",
     "CachePolicy",
-    "ContentText",
+    "ContentAudio",
     "ContentImage",
+    "ContentText",
+    "ContentVideo",
     "Content",
     "ChatMessage",
     "ChatMessageSystem",

inspect_ai/model/_call_tools.py

@@ -24,11 +24,17 @@ from typing import (
 from jsonschema import Draft7Validator
 from pydantic import BaseModel
 
-from inspect_ai._util.content import Content, ContentImage, ContentText
+from inspect_ai._util.content import (
+    Content,
+    ContentAudio,
+    ContentImage,
+    ContentText,
+    ContentVideo,
+)
 from inspect_ai._util.format import format_function_call
 from inspect_ai._util.text import truncate_string_to_bytes
 from inspect_ai._util.trace import trace_action
-from inspect_ai.model._trace import trace_tool_mesage
+from inspect_ai.model._conversation import conversation_tool_mesage
 from inspect_ai.tool import Tool, ToolCall, ToolError, ToolInfo
 from inspect_ai.tool._tool import ToolApprovalError, ToolParsingError
 from inspect_ai.tool._tool_call import ToolCallContent, ToolCallError
@@ -120,10 +126,14 @@ async def call_tools(
             # massage result, leave list[Content] alone, convert all other
             # types to string as that is what the model APIs accept
             truncated: tuple[int, int] | None = None
-            if isinstance(result, ContentText | ContentImage):
+            if isinstance(
+                result, ContentText | ContentImage | ContentAudio | ContentVideo
+            ):
                 content: str | list[Content] = [result]
             elif isinstance(result, list) and (
-                isinstance(result[0], ContentText | ContentImage)
+                isinstance(
+                    result[0], ContentText | ContentImage | ContentAudio | ContentVideo
+                )
             ):
                 content = result
             else:
@@ -163,6 +173,9 @@ async def call_tools(
         # call tools
         tool_messages: list[ChatMessageTool] = []
         for call in message.tool_calls:
+            # create the task
+            task = asyncio.create_task(call_tool_task(call))
+
             # create pending tool event and add it to the transcript
             event = ToolEvent(
                 id=call.id,
@@ -171,15 +184,44 @@ async def call_tools(
                 view=call.view,
                 pending=True,
             )
+            event.set_task(task)
             transcript()._event(event)
 
-            # execute the tool call
-            task = asyncio.create_task(call_tool_task(call))
-            tool_message, result_event = await task
+            # execute the tool call. if the operator cancelled the
+            # tool call then synthesize the appropriate message/event
+            try:
+                tool_message, result_event = await task
+            except asyncio.CancelledError:
+                if event.cancelled:
+                    tool_message = ChatMessageTool(
+                        content="",
+                        function=call.function,
+                        tool_call_id=call.id,
+                        error=ToolCallError(
+                            "timeout", "Command timed out before completing."
+                        ),
+                    )
+                    result_event = ToolEvent(
+                        id=call.id,
+                        function=call.function,
+                        arguments=call.arguments,
+                        result=tool_message.content,
+                        truncated=None,
+                        view=call.view,
+                        error=tool_message.error,
+                        events=[],
+                    )
+                    transcript().info(
+                        f"Tool call '{call.function}' was cancelled by operator."
+                    )
+                else:
+                    raise
+
+            # update return messages
             tool_messages.append(tool_message)
 
-            # trace if we are tracing
-            trace_tool_mesage(tool_message)
+            # print conversation if display is conversation
+            conversation_tool_mesage(tool_message)
 
             # update the event with the results
             event.set_result(
@@ -286,6 +328,10 @@ def tool_params(input: dict[str, Any], func: Callable[..., Any]) -> dict[str, An
     type_hints = get_type_hints(func)
     docstring = inspect.getdoc(func)
 
+    # if the function takes **kwargs: Any then just pass the tool arguments through
+    if "kwargs" in type_hints and type_hints["kwargs"] == Any:
+        return input
+
     # build params
     params: dict[str, Any] = {}
     for param_name, param in signature.parameters.items():
@@ -411,12 +457,13 @@ def truncate_tool_output(
     # truncate if required
     truncated = truncate_string_to_bytes(output, active_max_output)
     if truncated:
-        truncated_output = dedent(f"""
+        truncated_output = dedent("""
             The output of your call to {tool_name} was too long to be displayed.
             Here is a truncated version:
             <START_TOOL_OUTPUT>
-            {truncated.output}
-            <END_TOOL_OUTPUT>""")
+            {truncated_output}
+            <END_TOOL_OUTPUT>
+            """).format(tool_name=tool_name, truncated_output=truncated.output)
         return TruncatedToolOutput(
             truncated_output, truncated.original_bytes, active_max_output
         )

inspect_ai/model/_chat_message.py

@@ -59,10 +59,8 @@ class ChatMessageBase(BaseModel):
         if isinstance(self.content, str):
             self.content = text
         else:
-            all_images = [
-                content for content in self.content if content.type == "image"
-            ]
-            self.content = [ContentText(text=text)] + all_images
+            all_other = [content for content in self.content if content.type != "text"]
+            self.content = [ContentText(text=text)] + all_other
 
 
 class ChatMessageSystem(ChatMessageBase):

inspect_ai/model/_conversation.py

@@ -0,0 +1,61 @@
+from rich.console import RenderableType
+from rich.text import Text
+
+from inspect_ai._util.constants import NO_CONTENT
+from inspect_ai._util.rich import lines_display
+from inspect_ai._util.transcript import transcript_markdown
+from inspect_ai.util._conversation import conversation_panel
+from inspect_ai.util._display import display_type
+
+from ._chat_message import ChatMessage, ChatMessageAssistant, ChatMessageTool
+from ._render import messages_preceding_assistant, render_tool_calls
+
+MESSAGE_TITLE = "Message"
+
+
+def conversation_tool_mesage(message: ChatMessageTool) -> None:
+    if display_type() == "conversation":
+        # truncate output to 100 lines
+        output = (
+            message.error.message.strip() if message.error else message.text.strip()
+        )
+        if output:
+            content = lines_display(output, 100)
+
+            conversation_panel(
+                title=f"Tool Output: {message.function}",
+                content=content,
+            )
+
+
+def conversation_assistant_message(
+    input: list[ChatMessage], message: ChatMessageAssistant
+) -> None:
+    if display_type() == "conversation":
+        # print precding messages that aren't tool or assistant
+        for m in messages_preceding_assistant(input):
+            conversation_panel(
+                title=m.role.capitalize(),
+                content=transcript_markdown(m.text, escape=True),
+            )
+
+        # start with assistant content
+        content: list[RenderableType] = (
+            [transcript_markdown(message.text, escape=True)]
+            if message.text and message.text != NO_CONTENT
+            else []
+        )
+
+        # print tool calls
+        if message.tool_calls:
+            if content:
+                content.append(Text())
+            content.extend(render_tool_calls(message.tool_calls))
+
+        # print the assistant message
+        conversation_panel(title="Assistant", content=content)
+
+
+def conversation_assistant_error(error: Exception) -> None:
+    if display_type() == "conversation":
+        conversation_panel(title="Assistant", content=repr(error))

inspect_ai/model/_generate_config.py

@@ -58,14 +58,17 @@ class GenerateConfigArgs(TypedDict, total=False):
     """How many chat completion choices to generate for each input message. OpenAI, Grok, Google, and TogetherAI only."""
 
     logprobs: bool | None
-    """Return log probabilities of the output tokens. OpenAI, Google, Grok, TogetherAI, Huggingface, llama-cpp-python, and vLLM only."""
+    """Return log probabilities of the output tokens. OpenAI, Grok, TogetherAI, Huggingface, llama-cpp-python, and vLLM only."""
 
     top_logprobs: int | None
-    """Number of most likely tokens (0-20) to return at each token position, each with an associated log probability. OpenAI, Google, Grok, and Huggingface only."""
+    """Number of most likely tokens (0-20) to return at each token position, each with an associated log probability. OpenAI, Grok, and Huggingface only."""
 
     parallel_tool_calls: bool | None
     """Whether to enable parallel function calling during tool use (defaults to True). OpenAI and Groq only."""
 
+    internal_tools: bool | None
+    """Whether to automatically map tools to model internal implementations (e.g. 'computer' for anthropic)."""
+
     max_tool_output: int | None
     """Maximum tool output (in bytes). Defaults to 16 * 1024."""
 
@@ -128,14 +131,17 @@ class GenerateConfig(BaseModel):
     """How many chat completion choices to generate for each input message. OpenAI, Grok, Google, TogetherAI, and vLLM only."""
 
     logprobs: bool | None = Field(default=None)
-    """Return log probabilities of the output tokens. OpenAI, Google, Grok, TogetherAI, Huggingface, llama-cpp-python, and vLLM only."""
+    """Return log probabilities of the output tokens. OpenAI, Grok, TogetherAI, Huggingface, llama-cpp-python, and vLLM only."""
 
     top_logprobs: int | None = Field(default=None)
-    """Number of most likely tokens (0-20) to return at each token position, each with an associated log probability. OpenAI, Google, Grok, Huggingface, and vLLM only."""
+    """Number of most likely tokens (0-20) to return at each token position, each with an associated log probability. OpenAI, Grok, Huggingface, and vLLM only."""
 
     parallel_tool_calls: bool | None = Field(default=None)
     """Whether to enable parallel function calling during tool use (defaults to True). OpenAI and Groq only."""
 
+    internal_tools: bool | None = Field(default=None)
+    """Whether to automatically map tools to model internal implementations (e.g. 'computer' for anthropic)."""
+
     max_tool_output: int | None = Field(default=None)
     """Maximum tool output (in bytes). Defaults to 16 * 1024."""
 

inspect_ai/model/_model.py

@@ -33,6 +33,7 @@ from inspect_ai._util.trace import trace_action
 from inspect_ai.tool import Tool, ToolChoice, ToolFunction, ToolInfo
 from inspect_ai.tool._tool_def import ToolDef, tool_defs
 from inspect_ai.util import concurrency
+from inspect_ai.util._limit import SampleLimitExceededError
 
 from ._cache import CacheEntry, CachePolicy, cache_fetch, cache_store
 from ._call_tools import disable_parallel_tools, tool_call_view, tools_info
@@ -43,6 +44,7 @@ from ._chat_message import (
     ChatMessageTool,
     ChatMessageUser,
 )
+from ._conversation import conversation_assistant_error, conversation_assistant_message
 from ._generate_config import (
     GenerateConfig,
     active_generate_config,
@@ -50,7 +52,6 @@ from ._generate_config import (
 )
 from ._model_call import ModelCall
 from ._model_output import ModelOutput, ModelUsage
-from ._trace import trace_assistant_message
 
 logger = logging.getLogger(__name__)
 
@@ -116,7 +117,7 @@ class ModelAPI(abc.ABC):
         tools: list[ToolInfo],
         tool_choice: ToolChoice,
         config: GenerateConfig,
-    ) -> ModelOutput | tuple[ModelOutput, ModelCall]:
+    ) -> ModelOutput | tuple[ModelOutput | Exception, ModelCall]:
         """Generate output from the model.
 
         Args:
@@ -165,7 +166,7 @@ class ModelAPI(abc.ABC):
         return False
 
     def tool_result_images(self) -> bool:
-        """Tool results can containe images"""
+        """Tool results can contain images"""
         return False
 
 
@@ -222,11 +223,17 @@ class Model:
         Returns:
            ModelOutput
         """
+        # if we are the default model then enforce message limit if it
+        # exists (raise an exception if it is exceeded)
+        is_active_model = self == active_model()
+        if is_active_model:
+            handle_sample_message_limit(input)
+
         # base config for this model
         base_config = self.config
 
         # if we are the active_model then merge active generate config
-        if self == active_model():
+        if is_active_model:
             base_config = base_config.merge(active_generate_config())
 
         # merge passed config
@@ -296,6 +303,9 @@ class Model:
                 tools = []
             tool_choice = "none"
 
+        # apply any tool model_input handlers
+        input = resolve_tool_model_input(tdefs, input)
+
         # break tool image content out into user messages if the model doesn't
         # support tools returning images
         if not self.api.tool_result_images():
@@ -389,6 +399,17 @@ class Model:
                 output = result
                 call = None
 
+            # raise error
+            if isinstance(output, Exception):
+                complete(output, call)
+
+                # Wrap the error in a runtime error which will show the
+                # request which caused the error
+                error = repr(output)
+                request = json.dumps(call.request, indent=2) if call is not None else ""
+                error_message = f"{error}\n\nRequest:\n{request}"
+                raise RuntimeError(error_message)
+
             # update output with time elapsed
             output.time = time_elapsed
 
@@ -464,7 +485,7 @@ class Model:
         cache: Literal["read", "write"] | None,
         output: ModelOutput | None = None,
         call: ModelCall | None = None,
-    ) -> Callable[[ModelOutput, ModelCall | None], None]:
+    ) -> Callable[[ModelOutput | Exception, ModelCall | None], None]:
         from inspect_ai.log._transcript import ModelEvent, transcript
 
         # create event and add it to the transcript
@@ -484,13 +505,16 @@ class Model:
 
         # callable that can be used to update the interaction w/ output
         def complete(
-            updated_output: ModelOutput, updated_call: ModelCall | None
+            result: ModelOutput | Exception, updated_call: ModelCall | None
         ) -> None:
             # trace
-            trace_assistant_message(input, updated_output.choices[0].message)
+            if isinstance(result, ModelOutput):
+                conversation_assistant_message(input, result.choices[0].message)
+                event.output = result
+            else:
+                conversation_assistant_error(result)
+                event.error = repr(result)
 
-            # update event
-            event.output = updated_output
             event.call = updated_call
             event.pending = None
 
@@ -703,6 +727,40 @@ def simple_input_messages(
     return messages
 
 
+def resolve_tool_model_input(
+    tdefs: list[ToolDef], messages: list[ChatMessage]
+) -> list[ChatMessage]:
+    # filter on tooldefs that have a model input handler
+    tdefs = [tdef for tdef in tdefs if tdef.model_input is not None]
+
+    # bail if there are no handlers
+    if len(tdefs) == 0:
+        return messages
+
+    # don't mutate the original messages
+    messages = deepcopy(messages)
+
+    # extract tool messages
+    tool_messages = [
+        message for message in messages if isinstance(message, ChatMessageTool)
+    ]
+    # run model_input handlers over all tool_messages with the same function name
+    for tdef in tdefs:
+        assert tdef.model_input
+        # filter messages down to just this tool
+        tdef_tool_messages = [
+            message for message in tool_messages if message.function == tdef.name
+        ]
+        # call the function for each tool, passing the index, total, and content
+        for index, message in enumerate(tdef_tool_messages):
+            message.content = tdef.model_input(
+                index, len(tool_messages), message.content
+            )
+
+    # return modified messages
+    return messages
+
+
 def tool_result_images_as_user_message(
     messages: list[ChatMessage],
 ) -> list[ChatMessage]:
@@ -713,16 +771,21 @@ def tool_result_images_reducer(
     messages: list[ChatMessage],
     message: ChatMessage,
 ) -> list[ChatMessage]:
-    # append the message
-    messages.append(message)
-
     # if there are tool result images, pull them out into a ChatUserMessage
     if isinstance(message, ChatMessageTool) and isinstance(message.content, list):
+        tool_message = ChatMessageTool(
+            content=message.content.copy(),
+            tool_call_id=message.tool_call_id,
+            function=message.function,
+        )
+        assert isinstance(tool_message.content, list)
+        messages.append(tool_message)
+
         user_content: list[Content] = []
-        for i in range(0, len(message.content)):
-            if isinstance(message.content[i], ContentImage):
+        for i in range(0, len(tool_message.content)):
+            if isinstance(tool_message.content[i], ContentImage):
                 user_content.append(message.content[i])
-                message.content[i] = ContentText(
+                tool_message.content[i] = ContentText(
                     text="Image content is in the message below."
                 )
         if len(user_content) > 0:
@@ -730,6 +793,9 @@ def tool_result_images_reducer(
                 ChatMessageUser(content=user_content, tool_call_id=message.tool_call_id)
             )
 
+    else:
+        messages.append(message)
+
     # return messages
     return messages
 
@@ -813,6 +879,24 @@ def active_model() -> Model | None:
 active_model_context_var: ContextVar[Model] = ContextVar("active_model")
 
 
+def handle_sample_message_limit(input: str | list[ChatMessage]) -> None:
+    from inspect_ai.log._samples import (
+        active_sample_message_limit,
+        set_active_sample_total_messages,
+    )
+
+    total_messages = 1 if isinstance(input, str) else len(input)
+    message_limit = active_sample_message_limit()
+    if message_limit is not None:
+        if total_messages >= message_limit:
+            raise SampleLimitExceededError(
+                "message", value=total_messages, limit=message_limit
+            )
+
+    # set total messages
+    set_active_sample_total_messages(total_messages)
+
+
 def init_model_usage() -> None:
     model_usage_context_var.set({})
 
@@ -822,13 +906,28 @@ def init_sample_model_usage() -> None:
 
 
 def record_model_usage(model: str, usage: ModelUsage) -> None:
+    from inspect_ai.log._samples import (
+        active_sample_token_limit,
+        set_active_sample_total_tokens,
+    )
+
+    # record usage
     set_model_usage(model, usage, sample_model_usage_context_var.get(None))
     set_model_usage(model, usage, model_usage_context_var.get(None))
 
-    # update active sample
-    from inspect_ai.log._samples import set_active_sample_total_tokens
+    # compute total tokens
+    total_tokens = sample_total_tokens()
 
-    set_active_sample_total_tokens(sample_total_tokens())
+    # update active sample
+    set_active_sample_total_tokens(total_tokens)
+
+    # check for token limit overflow and raise
+    token_limit = active_sample_token_limit()
+    if token_limit is not None:
+        if total_tokens > token_limit:
+            raise SampleLimitExceededError(
+                "token", value=total_tokens, limit=token_limit
+            )
 
 
 def set_model_usage(

inspect_ai/model/_model_output.py

@@ -26,9 +26,14 @@ class ModelUsage(BaseModel):
 
 
 StopReason = Literal[
-    "stop", "max_tokens", "model_length", "tool_calls", "content_filter", "unknown"
+    "stop",
+    "max_tokens",
+    "model_length",
+    "tool_calls",
+    "content_filter",
+    "unknown",
 ]
-"""Reason that the model stopped generating."""
+"""Reason that the model stopped or failed to generate."""
 
 
 class TopLogprob(BaseModel):