inspect-ai 0.3.58__py3-none-any.whl → 0.3.60__py3-none-any.whl

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,7 +44,7 @@ from ._chat_message import (
     ChatMessageTool,
     ChatMessageUser,
 )
-from ._conversation import conversation_assistant_message
+from ._conversation import conversation_assistant_error, conversation_assistant_message
 from ._generate_config import (
     GenerateConfig,
     active_generate_config,
@@ -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
-            conversation_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):
@@ -209,3 +214,18 @@ class ModelOutput(BaseModel):
                 )
             ],
         )
+
+
+def as_stop_reason(reason: str | None) -> StopReason:
+    """Encode common reason strings into standard StopReason."""
+    match reason:
+        case "stop" | "eos":
+            return "stop"
+        case "length":
+            return "max_tokens"
+        case "tool_calls" | "function_call":
+            return "tool_calls"
+        case "content_filter" | "model_length" | "max_tokens":
+            return reason
+        case _:
+            return "unknown"

inspect_ai/model/_openai.py

@@ -0,0 +1,383 @@
+import json
+from typing import Literal
+
+from openai.types.chat import (
+    ChatCompletion,
+    ChatCompletionAssistantMessageParam,
+    ChatCompletionContentPartImageParam,
+    ChatCompletionContentPartInputAudioParam,
+    ChatCompletionContentPartParam,
+    ChatCompletionContentPartRefusalParam,
+    ChatCompletionContentPartTextParam,
+    ChatCompletionDeveloperMessageParam,
+    ChatCompletionMessage,
+    ChatCompletionMessageParam,
+    ChatCompletionMessageToolCall,
+    ChatCompletionMessageToolCallParam,
+    ChatCompletionNamedToolChoiceParam,
+    ChatCompletionSystemMessageParam,
+    ChatCompletionToolChoiceOptionParam,
+    ChatCompletionToolMessageParam,
+    ChatCompletionToolParam,
+    ChatCompletionUserMessageParam,
+)
+from openai.types.chat.chat_completion import Choice, ChoiceLogprobs
+from openai.types.chat.chat_completion_message_tool_call import Function
+from openai.types.completion_usage import CompletionUsage
+from openai.types.shared_params.function_definition import FunctionDefinition
+
+from inspect_ai._util.content import Content, ContentAudio, ContentImage, ContentText
+from inspect_ai._util.images import file_as_data_uri
+from inspect_ai._util.url import is_http_url
+from inspect_ai.model._call_tools import parse_tool_call
+from inspect_ai.model._model_output import ChatCompletionChoice, Logprobs
+from inspect_ai.tool import ToolCall, ToolChoice, ToolFunction, ToolInfo
+
+from ._chat_message import (
+    ChatMessage,
+    ChatMessageAssistant,
+    ChatMessageSystem,
+    ChatMessageTool,
+    ChatMessageUser,
+)
+from ._model_output import ModelUsage, StopReason, as_stop_reason
+
+
+def is_o1(name: str) -> bool:
+    return name.startswith("o1")
+
+
+def is_o1_full(name: str) -> bool:
+    return is_o1(name) and not is_o1_mini(name) and not is_o1_preview(name)
+
+
+def is_o1_mini(name: str) -> bool:
+    return name.startswith("o1-mini")
+
+
+def is_o1_preview(name: str) -> bool:
+    return name.startswith("o1-preview")
+
+
+def openai_chat_tool_call(tool_call: ToolCall) -> ChatCompletionMessageToolCall:
+    return ChatCompletionMessageToolCall(
+        type="function",
+        id=tool_call.id,
+        function=Function(
+            name=tool_call.function, arguments=json.dumps(tool_call.arguments)
+        ),
+    )
+
+
+def openai_chat_tool_call_param(
+    tool_call: ToolCall,
+) -> ChatCompletionMessageToolCallParam:
+    return ChatCompletionMessageToolCallParam(
+        id=tool_call.id,
+        function=dict(
+            name=tool_call.function, arguments=json.dumps(tool_call.arguments)
+        ),
+        type=tool_call.type,
+    )
+
+
+async def openai_chat_completion_part(
+    content: Content,
+) -> ChatCompletionContentPartParam:
+    if content.type == "text":
+        return ChatCompletionContentPartTextParam(type="text", text=content.text)
+    elif content.type == "image":
+        # API takes URL or base64 encoded file. If it's a remote file or
+        # data URL leave it alone, otherwise encode it
+        image_url = content.image
+        detail = content.detail
+
+        if not is_http_url(image_url):
+            image_url = await file_as_data_uri(image_url)
+
+        return ChatCompletionContentPartImageParam(
+            type="image_url",
+            image_url=dict(url=image_url, detail=detail),
+        )
+    elif content.type == "audio":
+        audio_data = await file_as_data_uri(content.audio)
+
+        return ChatCompletionContentPartInputAudioParam(
+            type="input_audio", input_audio=dict(data=audio_data, format=content.format)
+        )
+
+    else:
+        raise RuntimeError(
+            "Video content is not currently supported by Open AI chat models."
+        )
+
+
+async def openai_chat_message(
+    message: ChatMessage, model: str
+) -> ChatCompletionMessageParam:
+    if message.role == "system":
+        if is_o1(model):
+            return ChatCompletionDeveloperMessageParam(
+                role="developer", content=message.text
+            )
+        else:
+            return ChatCompletionSystemMessageParam(
+                role=message.role, content=message.text
+            )
+    elif message.role == "user":
+        return ChatCompletionUserMessageParam(
+            role=message.role,
+            content=(
+                message.content
+                if isinstance(message.content, str)
+                else [
+                    await openai_chat_completion_part(content)
+                    for content in message.content
+                ]
+            ),
+        )
+    elif message.role == "assistant":
+        if message.tool_calls:
+            return ChatCompletionAssistantMessageParam(
+                role=message.role,
+                content=message.text,
+                tool_calls=[
+                    openai_chat_tool_call_param(call) for call in message.tool_calls
+                ],
+            )
+        else:
+            return ChatCompletionAssistantMessageParam(
+                role=message.role, content=message.text
+            )
+    elif message.role == "tool":
+        return ChatCompletionToolMessageParam(
+            role=message.role,
+            content=(
+                f"Error: {message.error.message}" if message.error else message.text
+            ),
+            tool_call_id=str(message.tool_call_id),
+        )
+    else:
+        raise ValueError(f"Unexpected message role {message.role}")
+
+
+async def openai_chat_messages(
+    messages: list[ChatMessage], model: str
+) -> list[ChatCompletionMessageParam]:
+    return [await openai_chat_message(message, model) for message in messages]
+
+
+def openai_chat_choices(choices: list[ChatCompletionChoice]) -> list[Choice]:
+    oai_choices: list[Choice] = []
+
+    for index, choice in enumerate(choices):
+        if isinstance(choice.message.content, str):
+            content = choice.message.content
+        else:
+            content = "\n".join(
+                [c.text for c in choice.message.content if c.type == "text"]
+            )
+        if choice.message.tool_calls:
+            tool_calls = [openai_chat_tool_call(tc) for tc in choice.message.tool_calls]
+        else:
+            tool_calls = None
+        message = ChatCompletionMessage(
+            role="assistant", content=content, tool_calls=tool_calls
+        )
+        oai_choices.append(
+            Choice(
+                finish_reason=openai_finish_reason(choice.stop_reason),
+                index=index,
+                message=message,
+                logprobs=ChoiceLogprobs(**choice.logprobs.model_dump())
+                if choice.logprobs is not None
+                else None,
+            )
+        )
+
+    return oai_choices
+
+
+def openai_completion_usage(usage: ModelUsage) -> CompletionUsage:
+    return CompletionUsage(
+        completion_tokens=usage.output_tokens,
+        prompt_tokens=usage.input_tokens,
+        total_tokens=usage.total_tokens,
+    )
+
+
+def openai_finish_reason(
+    stop_reason: StopReason,
+) -> Literal["stop", "length", "tool_calls", "content_filter", "function_call"]:
+    match stop_reason:
+        case "stop" | "tool_calls" | "content_filter":
+            return stop_reason
+        case "model_length":
+            return "length"
+        case _:
+            return "stop"
+
+
+def openai_chat_tool_param(tool: ToolInfo) -> ChatCompletionToolParam:
+    function = FunctionDefinition(
+        name=tool.name,
+        description=tool.description,
+        parameters=tool.parameters.model_dump(exclude_none=True),
+    )
+    return ChatCompletionToolParam(type="function", function=function)
+
+
+def openai_chat_tools(tools: list[ToolInfo]) -> list[ChatCompletionToolParam]:
+    return [openai_chat_tool_param(tool) for tool in tools]
+
+
+def openai_chat_tool_choice(
+    tool_choice: ToolChoice,
+) -> ChatCompletionToolChoiceOptionParam:
+    if isinstance(tool_choice, ToolFunction):
+        return ChatCompletionNamedToolChoiceParam(
+            type="function", function=dict(name=tool_choice.name)
+        )
+    # openai supports 'any' via the 'required' keyword
+    elif tool_choice == "any":
+        return "required"
+    else:
+        return tool_choice
+
+
+def chat_tool_calls_from_openai(
+    message: ChatCompletionMessage, tools: list[ToolInfo]
+) -> list[ToolCall] | None:
+    if message.tool_calls:
+        return [
+            parse_tool_call(call.id, call.function.name, call.function.arguments, tools)
+            for call in message.tool_calls
+        ]
+    else:
+        return None
+
+
+def chat_messages_from_openai(
+    messages: list[ChatCompletionMessageParam],
+) -> list[ChatMessage]:
+    # track tool names by id
+    tool_names: dict[str, str] = {}
+
+    chat_messages: list[ChatMessage] = []
+
+    for message in messages:
+        if message["role"] == "system" or message["role"] == "developer":
+            sys_content = message["content"]
+            if isinstance(sys_content, str):
+                chat_messages.append(ChatMessageSystem(content=sys_content))
+            else:
+                chat_messages.append(
+                    ChatMessageSystem(
+                        content=[content_from_openai(c) for c in sys_content]
+                    )
+                )
+        elif message["role"] == "user":
+            user_content = message["content"]
+            if isinstance(user_content, str):
+                chat_messages.append(ChatMessageUser(content=user_content))
+            else:
+                chat_messages.append(
+                    ChatMessageUser(
+                        content=[content_from_openai(c) for c in user_content]
+                    )
+                )
+        elif message["role"] == "assistant":
+            # resolve content
+            asst_content = message["content"]
+            if isinstance(asst_content, str):
+                content: str | list[Content] = asst_content
+            elif asst_content is None:
+                content = message.get("refusal", None) or ""
+            else:
+                content = [content_from_openai(c) for c in asst_content]
+
+            # return message
+            if "tool_calls" in message:
+                tool_calls: list[ToolCall] = []
+                for tc in message["tool_calls"]:
+                    tool_calls.append(tool_call_from_openai(tc))
+                    tool_names[tc["id"]] = tc["function"]["name"]
+
+            else:
+                tool_calls = []
+            chat_messages.append(
+                ChatMessageAssistant(content=content, tool_calls=tool_calls or None)
+            )
+        elif message["role"] == "tool":
+            tool_content = message.get("content", None) or ""
+            if isinstance(tool_content, str):
+                content = tool_content
+            else:
+                content = [content_from_openai(c) for c in tool_content]
+            chat_messages.append(
+                ChatMessageTool(
+                    content=content,
+                    tool_call_id=message["tool_call_id"],
+                    function=tool_names.get(message["tool_call_id"], ""),
+                )
+            )
+        else:
+            raise ValueError(f"Unexpected message param type: {type(message)}")
+
+    return chat_messages
+
+
+def tool_call_from_openai(tool_call: ChatCompletionMessageToolCallParam) -> ToolCall:
+    return parse_tool_call(
+        tool_call["id"],
+        tool_call["function"]["name"],
+        tool_call["function"]["arguments"],
+    )
+
+
+def content_from_openai(
+    content: ChatCompletionContentPartParam | ChatCompletionContentPartRefusalParam,
+) -> Content:
+    if content["type"] == "text":
+        return ContentText(text=content["text"])
+    elif content["type"] == "image_url":
+        return ContentImage(
+            image=content["image_url"]["url"], detail=content["image_url"]["detail"]
+        )
+    elif content["type"] == "input_audio":
+        return ContentAudio(
+            audio=content["input_audio"]["data"],
+            format=content["input_audio"]["format"],
+        )
+    elif content["type"] == "refusal":
+        return ContentText(text=content["refusal"])
+
+
+def chat_message_assistant_from_openai(
+    message: ChatCompletionMessage, tools: list[ToolInfo]
+) -> ChatMessageAssistant:
+    refusal = getattr(message, "refusal", None)
+    return ChatMessageAssistant(
+        content=refusal or message.content or "",
+        source="generate",
+        tool_calls=chat_tool_calls_from_openai(message, tools),
+    )
+
+
+def chat_choices_from_openai(
+    response: ChatCompletion, tools: list[ToolInfo]
+) -> list[ChatCompletionChoice]:
+    choices = list(response.choices)
+    choices.sort(key=lambda c: c.index)
+    return [
+        ChatCompletionChoice(
+            message=chat_message_assistant_from_openai(choice.message, tools),
+            stop_reason=as_stop_reason(choice.finish_reason),
+            logprobs=(
+                Logprobs(**choice.logprobs.model_dump())
+                if choice.logprobs is not None
+                else None
+            ),
+        )
+        for choice in choices
+    ]