Merge pull request #106 from openai/add-attachment-metadata

Add a metadata field on `AttachmentBase` for integration-internal use cases

Author: jiwon-oai

chatkit/server.py

@@ -878,7 +878,11 @@ async def _paginate_thread_items_reverse(
             after = items.after
 
     def _serialize(self, obj: BaseModel) -> bytes:
-        return obj.model_dump_json(by_alias=True, exclude_none=True).encode("utf-8")
+        return obj.model_dump_json(
+            by_alias=True,
+            exclude_none=True,
+            context={"exclude_metadata": True},
+        ).encode("utf-8")
 
     def _to_thread_response(self, thread: ThreadMetadata | Thread) -> Thread:
         def is_hidden(item: ThreadItem) -> bool:

chatkit/types.py

@@ -3,7 +3,7 @@
 from datetime import datetime
 from typing import Any, Generic, Literal
 
-from pydantic import AnyUrl, BaseModel, Field
+from pydantic import AnyUrl, BaseModel, Field, SerializationInfo, model_serializer
 from typing_extensions import Annotated, TypeIs, TypeVar
 
 from chatkit.errors import ErrorCode
@@ -748,6 +748,19 @@ class AttachmentBase(BaseModel):
     Should be set to None after upload is complete or when using direct upload
     where uploading happens when creating the attachment object.
     """
+    metadata: dict[str, Any] | None = None
+    """
+    Integration-only metadata stored with the attachment. Ignored by ChatKit and not
+    returned in ChatKitServer responses. If you serialize attachments from a custom
+    direct-upload endpoint and want to omit this field, pass context={"exclude_metadata": True}.
+    """
+
+    @model_serializer(mode="wrap")
+    def _serialize(self, serializer, info: SerializationInfo):
+        data = serializer(self)
+        if isinstance(data, dict) and (info.context or {}).get("exclude_metadata"):
+            data.pop("metadata", None)
+        return data
 
 
 class FileAttachment(AttachmentBase):

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "openai-chatkit"
-version = "1.5.2"
+version = "1.5.3"
 description = "A ChatKit backend SDK."
 readme = "README.md"
 requires-python = ">=3.10"

tests/test_chatkit_server.py

@@ -106,6 +106,7 @@ async def create_attachment(
     ) -> Attachment:
         # Simple in-memory attachment creation
         id = f"atc_{len(self.files) + 1}"
+        metadata = {"source": "test"}
         if input.mime_type.startswith("image/"):
             attachment = ImageAttachment(
                 id=id,
@@ -117,6 +118,7 @@ async def create_attachment(
                     method="PUT",
                     headers={"X-My-Header": "my-value"},
                 ),
+                metadata=metadata,
             )
         else:
             attachment = FileAttachment(
@@ -128,6 +130,7 @@ async def create_attachment(
                     method="PUT",
                     headers={"X-My-Header": "my-value"},
                 ),
+                metadata=metadata,
             )
         self.files[attachment.id] = attachment
         return attachment
@@ -1100,6 +1103,10 @@ async def test_create_file():
         assert attachment.upload_descriptor.headers == {"X-My-Header": "my-value"}
         assert attachment.id in store.files
 
+        # The `metadata` field is excluded from serialization by default but still stored server-side
+        assert attachment.metadata is None
+        assert store.files[attachment.id].metadata == {"source": "test"}
+
 
 async def test_create_image_file():
     store = InMemoryFileStore()
@@ -1135,6 +1142,43 @@ async def test_create_image_file():
         assert attachment.id in store.files
 
 
+async def test_user_message_attachment_metadata_not_serialized():
+    attachment = FileAttachment(
+        id="file_with_metadata",
+        type="file",
+        mime_type="text/plain",
+        name="test.txt",
+        metadata={"source": "test"},
+    )
+
+    with make_server() as server:
+        await server.store.save_attachment(attachment, DEFAULT_CONTEXT)
+        events = await server.process_streaming(
+            ThreadsCreateReq(
+                params=ThreadCreateParams(
+                    input=UserMessageInput(
+                        content=[UserMessageTextContent(text="Message with file")],
+                        attachments=[attachment.id],
+                        inference_options=InferenceOptions(),
+                    )
+                )
+            )
+        )
+        thread = next(
+            event.thread for event in events if event.type == "thread.created"
+        )
+
+        list_result = await server.process_non_streaming(
+            ItemsListReq(params=ItemsListParams(thread_id=thread.id))
+        )
+        items = TypeAdapter(Page[ThreadItem]).validate_json(list_result.json)
+        user_items = [item for item in items.data if item.type == "user_message"]
+        assert user_items
+        attachments = user_items[0].attachments
+        assert attachments
+        assert all(att.metadata is None for att in attachments)
+
+
 async def test_create_file_without_filestore():
     with pytest.raises(RuntimeError):
         with make_server() as server:

tests/test_store.py

@@ -39,7 +39,15 @@ def make_thread_items() -> list[ThreadItem]:
     user_msg = UserMessageItem(
         id="msg_100000",
         content=[UserMessageTextContent(text="Hello!")],
-        attachments=[],
+        attachments=[
+            FileAttachment(
+                id="file_1",
+                type="file",
+                mime_type="text/plain",
+                name="test.txt",
+                metadata={"source": "test"},
+            )
+        ],
         inference_options=InferenceOptions(),
         thread_id="test_thread",
         created_at=now,
@@ -169,6 +177,7 @@ async def test_save_and_load_image(self):
             mime_type="image/png",
             name="test.png",
             preview_url=AnyUrl("https://example.com/test.png"),
+            metadata={"source": "test"},
         )
         await self.store.save_attachment(image, DEFAULT_CONTEXT)
         loaded = await self.store.load_attachment(image.id, DEFAULT_CONTEXT)