Model API: ChatMessage now includes an id field (#1437)

* add id to chat message

* ensure that message ids are preserved

* update changelog

* preserve id for tool

* update test to work with ids

---------

Co-authored-by: jjallaire <[email protected]>

CHANGELOG.md

@@ -6,6 +6,7 @@
 - New `inspect trace http` command which will show all HTTP requests for a run.
 - More consistent use of `max_retries` and `timeout` configuration options. These options now exclusively control Inspect's outer retry handler; model providers use their default behaviour for the inner request, which is typically 2-4 retries and a service-appropriate timeout.
 - Improved async implementation using AnyIO (can now optionally run Trio rather than asyncio as the [async backend](https://inspect.ai-safety-institute.org.uk/parallelism.html#async-backends)).
+- Model API: `ChatMessage` now includes an `id` field (defaults to auto-generated uuid).
 - Logging: Inspect no longer sets the global log level nor does it allow its own messages to propagate to the global handler (eliminating the possiblity of duplicate display). This should improve compatibility with applications that have their own custom logging configured. 
 - Tasks: For filesystem based tasks, no longer switch to the task file's directory during execution (directory switching still occurs during task loading). Specify `@task(chdir=True)` to preserve the previous behavior.
 - Bugfix: Fix issue with deserializing custom sandbox configuration objects.

src/inspect_ai/_eval/task/images.py

@@ -6,7 +6,7 @@
 from inspect_ai._util.images import file_as_data_uri
 from inspect_ai._util.url import is_data_uri
 from inspect_ai.dataset import Sample
-from inspect_ai.model import ChatMessage, ChatMessageUser
+from inspect_ai.model import ChatMessage
 from inspect_ai.solver import TaskState
 
 
@@ -66,27 +66,31 @@ def messages_without_base64_content(messages: list[ChatMessage]) -> list[ChatMes
 
 
 async def message_with_base64_content(message: ChatMessage) -> ChatMessage:
-    if isinstance(message, ChatMessageUser) and not isinstance(message.content, str):
-        return ChatMessageUser(
-            content=[
-                await chat_content_with_base64_content(content)
-                for content in message.content
-            ],
-            source=message.source,
+    if not isinstance(message.content, str):
+        return message.model_copy(
+            update=dict(
+                content=[
+                    await chat_content_with_base64_content(content)
+                    for content in message.content
+                ]
+            )
         )
+
     else:
         return message
 
 
 def message_without_base64_content(message: ChatMessage) -> ChatMessage:
-    if isinstance(message, ChatMessageUser) and not isinstance(message.content, str):
-        return ChatMessageUser(
-            content=[
-                chat_content_without_base64_content(content)
-                for content in message.content
-            ],
-            source=message.source,
+    if not isinstance(message.content, str):
+        return message.model_copy(
+            update=dict(
+                content=[
+                    chat_content_without_base64_content(content)
+                    for content in message.content
+                ]
+            )
         )
+
     else:
         return message
 

src/inspect_ai/_view/www/log-schema.json

@@ -199,6 +199,10 @@
     "ChatMessageAssistant": {
       "description": "Assistant chat message.",
       "properties": {
+        "id": {
+          "title": "Id",
+          "type": "string"
+        },
         "content": {
           "anyOf": [
             {
@@ -268,6 +272,7 @@
         }
       },
       "required": [
+        "id",
         "content",
         "source",
         "role",
@@ -280,6 +285,10 @@
     "ChatMessageSystem": {
       "description": "System chat message.",
       "properties": {
+        "id": {
+          "title": "Id",
+          "type": "string"
+        },
         "content": {
           "anyOf": [
             {
@@ -334,6 +343,7 @@
         }
       },
       "required": [
+        "id",
         "content",
         "source",
         "role"
@@ -345,6 +355,10 @@
     "ChatMessageTool": {
       "description": "Tool chat message.",
       "properties": {
+        "id": {
+          "title": "Id",
+          "type": "string"
+        },
         "content": {
           "anyOf": [
             {
@@ -434,6 +448,7 @@
         }
       },
       "required": [
+        "id",
         "content",
         "source",
         "role",
@@ -448,6 +463,10 @@
     "ChatMessageUser": {
       "description": "User chat message.",
       "properties": {
+        "id": {
+          "title": "Id",
+          "type": "string"
+        },
         "content": {
           "anyOf": [
             {
@@ -517,6 +536,7 @@
         }
       },
       "required": [
+        "id",
         "content",
         "source",
         "role",

src/inspect_ai/_view/www/src/types/log.d.ts

@@ -134,6 +134,7 @@ export type Input =
       | ChatMessageAssistant
       | ChatMessageTool
     )[];
+export type Id1 = string;
 export type Content =
   | string
   | (
@@ -160,6 +161,7 @@ export type Video = string;
 export type Format1 = "mp4" | "mpeg" | "mov";
 export type Source = ("input" | "generate") | null;
 export type Role = "system";
+export type Id2 = string;
 export type Content1 =
   | string
   | (
@@ -172,6 +174,7 @@ export type Content1 =
 export type Source1 = ("input" | "generate") | null;
 export type Role1 = "user";
 export type ToolCallId = string[] | null;
+export type Id3 = string;
 export type Content2 =
   | string
   | (
@@ -184,13 +187,14 @@ export type Content2 =
 export type Source2 = ("input" | "generate") | null;
 export type Role2 = "assistant";
 export type ToolCalls = ToolCall[] | null;
-export type Id1 = string;
+export type Id4 = string;
 export type Function = string;
 export type Type7 = "function";
 export type ParseError = string | null;
 export type Title = string | null;
 export type Format2 = "text" | "markdown";
 export type Content3 = string;
+export type Id5 = string;
 export type Content4 =
   | string
   | (
@@ -273,7 +277,7 @@ export type Input1 =
     )[];
 export type Choices2 = string[] | null;
 export type Target1 = string | string[];
-export type Id2 = number | string | null;
+export type Id6 = number | string | null;
 export type Metadata8 = {} | null;
 export type Files1 = {
   [k: string]: string;
@@ -360,7 +364,7 @@ export type WorkingStart6 = number;
 export type Pending6 = boolean | null;
 export type Event6 = "tool";
 export type Type12 = "function";
-export type Id3 = string;
+export type Id7 = string;
 export type Function2 = string;
 export type Result1 =
   | string
@@ -805,6 +809,7 @@ export interface EvalSample {
  * System chat message.
  */
 export interface ChatMessageSystem {
+  id: Id1;
   content: Content;
   source: Source;
   role: Role;
@@ -855,6 +860,7 @@ export interface ContentVideo {
  * User chat message.
  */
 export interface ChatMessageUser {
+  id: Id2;
   content: Content1;
   source: Source1;
   role: Role1;
@@ -864,13 +870,14 @@ export interface ChatMessageUser {
  * Assistant chat message.
  */
 export interface ChatMessageAssistant {
+  id: Id3;
   content: Content2;
   source: Source2;
   role: Role2;
   tool_calls: ToolCalls;
 }
 export interface ToolCall {
-  id: Id1;
+  id: Id4;
   function: Function;
   arguments: Arguments;
   type: Type7;
@@ -890,6 +897,7 @@ export interface ToolCallContent {
  * Tool chat message.
  */
 export interface ChatMessageTool {
+  id: Id5;
   content: Content4;
   source: Source3;
   role: Role3;
@@ -972,7 +980,7 @@ export interface Sample {
   input: Input1;
   choices: Choices2;
   target: Target1;
-  id: Id2;
+  id: Id6;
   metadata: Metadata8;
   sandbox: SandboxEnvironmentSpec | null;
   files: Files1;
@@ -1176,7 +1184,7 @@ export interface ToolEvent {
   pending: Pending6;
   event: Event6;
   type: Type12;
-  id: Id3;
+  id: Id7;
   function: Function2;
   arguments: Arguments1;
   view: ToolCallContent | null;

src/inspect_ai/dataset/_sources/util.py

@@ -63,12 +63,13 @@ def message_with_resolved_content(
     message: ChatMessage, resolver: Callable[[str], str]
 ) -> ChatMessage:
     if isinstance(message, ChatMessageUser) and not isinstance(message.content, str):
-        return ChatMessageUser(
-            content=[
-                chat_content_with_resolved_content(content, resolver)
-                for content in message.content
-            ],
-            source=message.source,
+        return message.model_copy(
+            update=dict(
+                content=[
+                    chat_content_with_resolved_content(content, resolver)
+                    for content in message.content
+                ],
+            )
         )
     else:
         return message

src/inspect_ai/model/_chat_message.py

@@ -2,6 +2,7 @@
 from typing import Any, Literal, Type, Union
 
 from pydantic import BaseModel, Field, model_validator
+from shortuuid import uuid
 
 from inspect_ai._util.content import Content, ContentReasoning, ContentText
 from inspect_ai.tool import ToolCall
@@ -15,6 +16,9 @@
 class ChatMessageBase(BaseModel):
     """Base class for chat messages."""
 
+    id: str = Field(default_factory=uuid)
+    """Unique identifer for message."""
+
     content: str | list[Content]
     """Content (simple string or list of content objects)"""
 

src/inspect_ai/model/_model.py

@@ -1094,6 +1094,7 @@ def tool_result_images_reducer(
             messages
             + [
                 ChatMessageTool(
+                    id=message.id,
                     content=edited_tool_message_content,
                     tool_call_id=message.tool_call_id,
                     function=message.function,
@@ -1200,13 +1201,13 @@ def combine_messages(
     a: ChatMessage, b: ChatMessage, message_type: Type[ChatMessage]
 ) -> ChatMessage:
     if isinstance(a.content, str) and isinstance(b.content, str):
-        return message_type(content=f"{a.content}\n{b.content}")
+        return message_type(id=a.id, content=f"{a.content}\n{b.content}")
     elif isinstance(a.content, list) and isinstance(b.content, list):
-        return message_type(content=a.content + b.content)
+        return message_type(id=a.id, content=a.content + b.content)
     elif isinstance(a.content, str) and isinstance(b.content, list):
-        return message_type(content=[ContentText(text=a.content), *b.content])
+        return message_type(id=a.id, content=[ContentText(text=a.content), *b.content])
     elif isinstance(a.content, list) and isinstance(b.content, str):
-        return message_type(content=a.content + [ContentText(text=b.content)])
+        return message_type(id=a.id, content=a.content + [ContentText(text=b.content)])
     else:
         raise TypeError(
             f"Cannot combine messages with invalid content types: {a.content!r}, {b.content!r}"

src/inspect_ai/model/_providers/openai_o1.py

@@ -107,7 +107,7 @@ def chat_messages(
 ) -> list[ChatCompletionMessageParam]:
     # o1 does not allow system messages so convert system -> user
     messages: list[ChatMessage] = [
-        ChatMessageUser(content=message.content)
+        ChatMessageUser(id=message.id, content=message.content)
         if message.role == "system"
         else message
         for message in input

tests/model/test_tool_result_images_as_user_message.py

@@ -20,7 +20,8 @@ def test_multiple_tool_responses_remain_adjacent():
             _modified_image_response_message(tool_a),
             _modified_image_response_message(tool_b),
             ChatMessageUser(
-                content=tool_a.content + tool_b.content, tool_call_id=["a", "b"]
+                content=tool_a.content + tool_b.content,
+                tool_call_id=["a", "b"],
             ),
         ],
     )
@@ -43,15 +44,20 @@ def test_multiple_tool_responses_remain_adjacent_when_not_at_end_of_list():
             _modified_image_response_message(tool_a),
             _modified_image_response_message(tool_b),
             ChatMessageUser(
-                content=tool_a.content + tool_b.content, tool_call_id=["a", "b"]
+                content=tool_a.content + tool_b.content,
+                tool_call_id=["a", "b"],
             ),
             user,
         ],
     )
 
 
 def execute_and_assert(input_messages: list[ChatMessage], expected: list[ChatMessage]):
-    assert tool_result_images_as_user_message(input_messages) == expected
+    # transform the messages
+    transformed = tool_result_images_as_user_message(input_messages)
+
+    # compare based on content (as id can't be known in advance)
+    assert all([t.content == e.content] for t, e in zip(transformed, expected))
 
 
 def _modified_image_response_message(message: ChatMessageTool) -> ChatMessageTool: