use new type of hidden context item; add stream options

Author: jiwon-oai

chatkit/agents.py

@@ -56,6 +56,7 @@
     EndOfTurnItem,
     FileSource,
     HiddenContextItem,
+    SDKHiddenContextItem,
     Task,
     TaskItem,
     ThoughtTask,
@@ -691,9 +692,6 @@ async def hidden_context_to_input(
         """
         Convert a HiddenContextItem into input item(s) to send to the model.
         Required to override when HiddenContextItems with non-string content are used.
-
-        ChatKitServer may save HiddenContextItems with string content; make sure your override
-        can also handle HiddenContextItems with string content.
         """
         if not isinstance(item.content, str):
             raise NotImplementedError(
@@ -715,6 +713,29 @@ async def hidden_context_to_input(
             role="user",
         )
 
+    async def sdk_hidden_context_to_input(
+        self, item: SDKHiddenContextItem
+    ) -> TResponseInputItem | list[TResponseInputItem] | None:
+        """
+        Convert a SDKHiddenContextItem into input item to send to the model.
+        This is used by the ChatKit Python SDK for storing additional context
+        for internal operations; you shouldn't need to override this.
+        """
+        text = (
+            "Hidden ChatKit SDK context for the agent (not shown to the user):\n"
+            f"<SDKHiddenContext>\n{item.content}\n</SDKHiddenContext>"
+        )
+        return Message(
+            type="message",
+            content=[
+                ResponseInputTextParam(
+                    type="input_text",
+                    text=text,
+                )
+            ],
+            role="user",
+        )
+
     async def task_to_input(
         self, item: TaskItem
     ) -> TResponseInputItem | list[TResponseInputItem] | None:
@@ -951,6 +972,9 @@ async def _thread_item_to_input_item(
             case HiddenContextItem():
                 out = await self.hidden_context_to_input(item) or []
                 return out if isinstance(out, list) else [out]
+            case SDKHiddenContextItem():
+                out = await self.sdk_hidden_context_to_input(item) or []
+                return out if isinstance(out, list) else [out]
             case _:
                 assert_never(item)
 

chatkit/server.py

@@ -45,7 +45,10 @@
     ItemsListReq,
     NonStreamingReq,
     Page,
+    SDKHiddenContextItem,
     StreamingReq,
+    StreamOptions,
+    StreamOptionsEvent,
     Thread,
     ThreadCreatedEvent,
     ThreadItem,
@@ -73,6 +76,8 @@
     WidgetRootUpdated,
     WidgetStreamingTextValueDelta,
     WorkflowItem,
+    WorkflowTaskAdded,
+    WorkflowTaskUpdated,
     is_streaming_req,
 )
 from .version import __version__
@@ -315,24 +320,55 @@ def action(
             "See https://github.com/openai/chatkit-python/blob/main/docs/widgets.md#widget-actions"
         )
 
+    def get_stream_options(
+        self, thread: ThreadMetadata, context: TContext
+    ) -> StreamOptions:
+        """
+        Return stream-level runtime options. Allows the user to cancel the stream by default.
+        Override this method to customize behavior.
+        """
+        return StreamOptions(allow_cancel=True)
+
     async def handle_stream_cancelled(
         self,
         thread: ThreadMetadata,
         pending_items: list[ThreadItem],
         context: TContext,
     ):
         """Perform custom cleanup / stop inference when a stream is cancelled.
+        Updates you make here will not be reflected in the UI until a reload.
+
+        The default implementation persists any non-empty pending assistant messages
+        to the thread but does not auto-save pending widget items or workflow items.
 
         Args:
             thread: The thread that was being processed.
             pending_items: Items that were not done streaming at cancellation time.
-                By default, already-streamed assistant messages, widgets, and workflows are
-                saved to the store during error handling prior to this method being called.
-                If you want to remove them from the thread, you can do so here.
-                (Updates you make here will not be reflected in the UI until a reload.)
             context: Arbitrary per-request context provided by the caller.
         """
-        pass
+        pending_assistant_message_items: list[AssistantMessageItem] = [
+            item for item in pending_items if isinstance(item, AssistantMessageItem)
+        ]
+        for item in pending_assistant_message_items:
+            is_empty = len(item.content) == 0 or all(
+                (not content.text.strip()) for content in item.content
+            )
+            if not is_empty:
+                await self.store.add_thread_item(thread.id, item, context=context)
+
+        # Add a hidden context item to the thread to indicate that the stream was cancelled.
+        # Otherwise, depending on the timing of the cancellation, subsequent responses may
+        # attempt to continue the cancelled response.
+        await self.store.add_thread_item(
+            thread.id,
+            SDKHiddenContextItem(
+                thread_id=thread.id,
+                created_at=datetime.now(),
+                id=self.store.generate_item_id("sdk_hidden_context", thread, context),
+                content="The user cancelled the stream.",
+            ),
+            context=context,
+        )
 
     async def process(
         self, request: str | bytes | bytearray, context: TContext
@@ -633,6 +669,11 @@ async def _process_events(
     ) -> AsyncIterator[ThreadStreamEvent]:
         await asyncio.sleep(0)  # allow the response to start streaming
 
+        # Send initial stream options
+        yield StreamOptionsEvent(
+            stream_options=self.get_stream_options(thread, context)
+        )
+
         last_thread = thread.model_copy(deep=True)
 
         # Keep track of items that were streamed but not yet saved
@@ -662,12 +703,10 @@ async def _process_events(
                             )
                             pending_items.pop(event.item.id, None)
                         case ThreadItemUpdatedEvent():
-                            # Keep the pending assistant message item up to date
-                            # so that we can persist already-streamed partial content
-                            # if the stream is cancelled.
-                            self._update_pending_assistant_message_items(
-                                pending_items, event
-                            )
+                            # Keep pending assistant message and workflow items up to date
+                            # so that we have a reference to the latest version of these pending items
+                            # when the stream is cancelled.
+                            self._update_pending_items(pending_items, event)
 
                     # special case - don't send hidden context items back to the client
                     should_swallow_event = isinstance(
@@ -690,10 +729,6 @@ async def _process_events(
                     await self.store.save_thread(thread, context=context)
                     yield ThreadUpdatedEvent(thread=self._to_thread_response(thread))
         except asyncio.CancelledError:
-            # When a stream is cancelled, whether it's a deliberate stop request or due to a network issue,
-            # save already-streamed items to the thread.
-            await self._persist_cancelled_stream_state(thread, pending_items, context)
-            # Allow custom cleanup.
             await self.handle_stream_cancelled(
                 thread, list(pending_items.values()), context
             )
@@ -721,43 +756,6 @@ async def _process_events(
             await self.store.save_thread(thread, context=context)
             yield ThreadUpdatedEvent(thread=self._to_thread_response(thread))
 
-    async def _persist_cancelled_stream_state(
-        self,
-        thread: ThreadMetadata,
-        pending_items: dict[str, ThreadItem],
-        context: TContext,
-    ):
-        # Persist any streamed items that the UI should keep when cancellation happens mid-stream.
-        for item in pending_items.values():
-            if isinstance(
-                item, (AssistantMessageItem, WidgetItem, WorkflowItem)
-            ) and not self._is_streamed_item_empty(item):
-                await self.store.add_thread_item(thread.id, item, context=context)
-
-        await self.store.add_thread_item(
-            thread.id,
-            HiddenContextItem(
-                thread_id=thread.id,
-                created_at=datetime.now(),
-                id=self.store.generate_item_id("hidden_context", thread, context),
-                content="SYSTEM: The user cancelled the stream.",
-            ),
-            context=context,
-        )
-
-    def _is_streamed_item_empty(
-        self, item: AssistantMessageItem | WorkflowItem | WidgetItem
-    ) -> bool:
-        if isinstance(item, AssistantMessageItem):
-            return len(item.content) == 0 or all(
-                (not content.text.strip()) for content in item.content
-            )
-        if isinstance(item, WorkflowItem):
-            return len(item.workflow.tasks) == 0 and item.workflow.summary is None
-
-        # Assume all WidgetItems are not empty
-        return False
-
     def _apply_assistant_message_update(
         self,
         item: AssistantMessageItem,
@@ -788,27 +786,38 @@ def _apply_assistant_message_update(
                 updated.content[update.content_index] = update.content
         return updated
 
-    def _update_pending_assistant_message_items(
+    def _update_pending_items(
         self,
         pending_items: dict[str, ThreadItem],
         event: ThreadItemUpdatedEvent,
     ):
-        if not isinstance(
-            event.update,
-            (
-                AssistantMessageContentPartAdded,
-                AssistantMessageContentPartTextDelta,
-                AssistantMessageContentPartAnnotationAdded,
-                AssistantMessageContentPartDone,
-            ),
-        ):
-            return
-
         updated_item = pending_items.get(event.item_id)
-        if updated_item and isinstance(updated_item, AssistantMessageItem):
-            pending_items[updated_item.id] = self._apply_assistant_message_update(
-                updated_item, event.update
-            )
+        update = event.update
+        match updated_item:
+            case AssistantMessageItem():
+                if isinstance(
+                    update,
+                    (
+                        AssistantMessageContentPartAdded,
+                        AssistantMessageContentPartTextDelta,
+                        AssistantMessageContentPartAnnotationAdded,
+                        AssistantMessageContentPartDone,
+                    ),
+                ):
+                    pending_items[updated_item.id] = (
+                        self._apply_assistant_message_update(updated_item, update)
+                    )
+            case WorkflowItem():
+                if isinstance(update, (WorkflowTaskUpdated, WorkflowTaskAdded)):
+                    match update:
+                        case WorkflowTaskUpdated():
+                            updated_item.workflow.tasks[update.task_index] = update.task
+                        case WorkflowTaskAdded():
+                            updated_item.workflow.tasks.append(update.task)
+
+                    pending_items[updated_item.id] = updated_item
+            case _:
+                pass
 
     async def _build_user_message_item(
         self, input: UserMessageInput, thread: ThreadMetadata, context: TContext

chatkit/store.py

@@ -15,7 +15,13 @@
 TContext = TypeVar("TContext", default=Any)
 
 StoreItemType = Literal[
-    "thread", "message", "tool_call", "task", "workflow", "attachment", "hidden_context"
+    "thread",
+    "message",
+    "tool_call",
+    "task",
+    "workflow",
+    "attachment",
+    "sdk_hidden_context",
 ]
 
 
@@ -26,7 +32,7 @@
     "workflow": "wf",
     "task": "tsk",
     "attachment": "atc",
-    "hidden_context": "hcx",
+    "sdk_hidden_context": "shcx",
 }
 
 

chatkit/types.py

@@ -317,6 +317,20 @@ class ThreadItemReplacedEvent(BaseModel):
     item: ThreadItem
 
 
+class StreamOptions(BaseModel):
+    """Settings that control runtime stream behavior."""
+
+    allow_cancel: bool
+    """Allow the client to request cancellation mid-stream."""
+
+
+class StreamOptionsEvent(BaseModel):
+    """Event emitted to set stream options at runtime."""
+
+    type: Literal["stream_options"] = "stream_options"
+    stream_options: StreamOptions
+
+
 class ProgressUpdateEvent(BaseModel):
     """Event providing incremental progress from the assistant."""
 
@@ -354,6 +368,7 @@ class NoticeEvent(BaseModel):
     | ThreadItemUpdated
     | ThreadItemRemovedEvent
     | ThreadItemReplacedEvent
+    | StreamOptionsEvent
     | ProgressUpdateEvent
     | ErrorEvent
     | NoticeEvent,
@@ -576,12 +591,25 @@ class EndOfTurnItem(ThreadItemBase):
 
 
 class HiddenContextItem(ThreadItemBase):
-    """HiddenContext is never sent to the client. It's not officially part of ChatKit. It is only used internally to store additional context in a specific place in the thread."""
+    """
+    HiddenContext is never sent to the client. It's not officially part of ChatKit.js.
+    It is only used internally to store additional context in a specific place in the thread.
+    """
 
     type: Literal["hidden_context_item"] = "hidden_context_item"
     content: Any
 
 
+class SDKHiddenContextItem(ThreadItemBase):
+    """
+    Hidden context that is used by the ChatKit Python SDK for storing additional context
+    for internal operations.
+    """
+
+    type: Literal["sdk_hidden_context"] = "sdk_hidden_context"
+    content: str
+
+
 ThreadItem = Annotated[
     UserMessageItem
     | AssistantMessageItem
@@ -590,6 +618,7 @@ class HiddenContextItem(ThreadItemBase):
     | WorkflowItem
     | TaskItem
     | HiddenContextItem
+    | SDKHiddenContextItem
     | EndOfTurnItem,
     Field(discriminator="type"),
 ]