feat: add plan_update chunk for chat streaming

slack_sdk/models/messages/chunk.py

@@ -0,0 +1,180 @@
+import logging
+from typing import Any, Dict, Optional, Sequence, Set, Union
+
+from slack_sdk.errors import SlackObjectFormationError
+from slack_sdk.models import show_unknown_key_warning
+from slack_sdk.models.basic_objects import JsonObject
+
+
+class Chunk(JsonObject):
+    """
+    Chunk for streaming messages.
+
+    https://docs.slack.dev/messaging/sending-and-scheduling-messages#text-streaming
+    """
+
+    attributes = {"type"}
+    logger = logging.getLogger(__name__)
+
+    def __init__(
+        self,
+        *,
+        type: Optional[str] = None,
+    ):
+        self.type = type
+
+    @classmethod
+    def parse(cls, chunk: Union[Dict, "Chunk"]) -> Optional["Chunk"]:
+        if chunk is None:
+            return None
+        elif isinstance(chunk, Chunk):
+            return chunk
+        else:
+            if "type" in chunk:
+                type = chunk["type"]
+                if type == MarkdownTextChunk.type:
+                    return MarkdownTextChunk(**chunk)
+                elif type == PlanUpdateChunk.type:
+                    return PlanUpdateChunk(**chunk)
+                elif type == TaskUpdateChunk.type:
+                    return TaskUpdateChunk(**chunk)
+                else:
+                    cls.logger.warning(f"Unknown chunk detected and skipped ({chunk})")
+                    return None
+            else:
+                cls.logger.warning(f"Unknown chunk detected and skipped ({chunk})")
+                return None
+
+
+class MarkdownTextChunk(Chunk):
+    type = "markdown_text"
+
+    @property
+    def attributes(self) -> Set[str]:  # type: ignore[override]
+        return super().attributes.union({"text"})
+
+    def __init__(
+        self,
+        *,
+        text: str,
+        **others: Dict,
+    ):
+        """Used for streaming text content with markdown formatting support.
+
+        https://docs.slack.dev/messaging/sending-and-scheduling-messages#text-streaming
+        """
+        super().__init__(type=self.type)
+        show_unknown_key_warning(self, others)
+
+        self.text = text
+
+
+class PlanUpdateChunk(Chunk):
+    type = "plan_update"
+
+    @property
+    def attributes(self) -> Set[str]:  # type: ignore[override]
+        return super().attributes.union({"title"})
+
+    def __init__(
+        self,
+        *,
+        title: str,
+        **others: Dict,
+    ):
+        """An updated title of plans for task and tool calls.
+
+        https://docs.slack.dev/messaging/sending-and-scheduling-messages#text-streaming
+        """
+        super().__init__(type=self.type)
+        show_unknown_key_warning(self, others)
+
+        self.title = title
+
+
+class URLSource(JsonObject):
+    type = "url"
+
+    @property
+    def attributes(self) -> Set[str]:
+        return super().attributes.union(
+            {
+                "url",
+                "text",
+                "icon_url",
+            }
+        )
+
+    def __init__(
+        self,
+        *,
+        url: str,
+        text: str,
+        icon_url: Optional[str] = None,
+        **others: Dict,
+    ):
+        show_unknown_key_warning(self, others)
+        self._url = url
+        self._text = text
+        self._icon_url = icon_url
+
+    def to_dict(self) -> Dict[str, Any]:
+        self.validate_json()
+        json: Dict[str, Union[str, Dict]] = {
+            "type": self.type,
+            "url": self._url,
+            "text": self._text,
+        }
+        if self._icon_url:
+            json["icon_url"] = self._icon_url
+        return json
+
+
+class TaskUpdateChunk(Chunk):
+    type = "task_update"
+
+    @property
+    def attributes(self) -> Set[str]:  # type: ignore[override]
+        return super().attributes.union(
+            {
+                "id",
+                "title",
+                "status",
+                "details",
+                "output",
+                "sources",
+            }
+        )
+
+    def __init__(
+        self,
+        *,
+        id: str,
+        title: str,
+        status: str,  # "pending", "in_progress", "complete", "error"
+        details: Optional[str] = None,
+        output: Optional[str] = None,
+        sources: Optional[Sequence[Union[Dict, URLSource]]] = None,
+        **others: Dict,
+    ):
+        """Used for displaying tool execution progress in a timeline-style UI.
+
+        https://docs.slack.dev/messaging/sending-and-scheduling-messages#text-streaming
+        """
+        super().__init__(type=self.type)
+        show_unknown_key_warning(self, others)
+
+        self.id = id
+        self.title = title
+        self.status = status
+        self.details = details
+        self.output = output
+        if sources is not None:
+            self.sources = []
+            for src in sources:
+                if isinstance(src, Dict):
+                    self.sources.append(src)
+                elif isinstance(src, URLSource):
+                    self.sources.append(src.to_dict())
+                else:
+                    raise SlackObjectFormationError(f"Unsupported type for source in task update chunk: {type(src)}")

slack_sdk/web/async_chat_stream.py

@@ -10,10 +10,11 @@
 
 import json
 import logging
-from typing import TYPE_CHECKING, Dict, Optional, Sequence, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Union
 
 import slack_sdk.errors as e
 from slack_sdk.models.blocks.blocks import Block
+from slack_sdk.models.messages.chunk import Chunk, MarkdownTextChunk
 from slack_sdk.models.metadata import Metadata
 from slack_sdk.web.async_slack_response import AsyncSlackResponse
 
@@ -75,7 +76,8 @@ def __init__(
     async def append(
         self,
         *,
-        markdown_text: str,
+        markdown_text: Optional[str] = None,
+        chunks: Optional[Sequence[Union[Dict, Chunk]]] = None,
         **kwargs,
     ) -> Optional[AsyncSlackResponse]:
         """Append to the stream.
@@ -84,6 +86,7 @@ async def append(
         is stopped this method cannot be called.
 
         Args:
+            chunks: An array of streaming chunks that can contain either markdown text or task updates.
             markdown_text: Accepts message text formatted in markdown. Limit this field to 12,000 characters. This text is
               what will be appended to the message received so far.
             **kwargs: Additional arguments passed to the underlying API calls.
@@ -111,9 +114,10 @@ async def append(
             raise e.SlackRequestError(f"Cannot append to stream: stream state is {self._state}")
         if kwargs.get("token"):
             self._token = kwargs.pop("token")
-        self._buffer += markdown_text
-        if len(self._buffer) >= self._buffer_size:
-            return await self._flush_buffer(**kwargs)
+        if markdown_text is not None:
+            self._buffer += markdown_text
+        if len(self._buffer) >= self._buffer_size or chunks is not None:
+            return await self._flush_buffer(chunks=chunks, **kwargs)
         details = {
             "buffer_length": len(self._buffer),
             "buffer_size": self._buffer_size,
@@ -129,6 +133,7 @@ async def stop(
         self,
         *,
         markdown_text: Optional[str] = None,
+        chunks: Optional[Sequence[Union[Dict, Chunk]]] = None,
         blocks: Optional[Union[str, Sequence[Union[Dict, Block]]]] = None,
         metadata: Optional[Union[Dict, Metadata]] = None,
         **kwargs,
@@ -137,6 +142,7 @@ async def stop(
 
         Args:
             blocks: A list of blocks that will be rendered at the bottom of the finalized message.
+            chunks: An array of streaming chunks that can contain either markdown text or task updates.
             markdown_text: Accepts message text formatted in markdown. Limit this field to 12,000 characters. This text is
               what will be appended to the message received so far.
             metadata: JSON object with event_type and event_payload fields, presented as a URL-encoded string. Metadata you
@@ -177,26 +183,36 @@ async def stop(
                 raise e.SlackRequestError("Failed to stop stream: stream not started")
             self._stream_ts = str(response["ts"])
             self._state = "in_progress"
+        flushings: List[Union[Dict, Chunk]] = []
+        if len(self._buffer) != 0:
+            flushings.append(MarkdownTextChunk(text=self._buffer))
+        if chunks is not None:
+            flushings.extend(chunks)
         response = await self._client.chat_stopStream(
             token=self._token,
             channel=self._stream_args["channel"],
             ts=self._stream_ts,
             blocks=blocks,
-            markdown_text=self._buffer,
+            chunks=flushings,
             metadata=metadata,
             **kwargs,
         )
         self._state = "completed"
         return response
 
-    async def _flush_buffer(self, **kwargs) -> AsyncSlackResponse:
-        """Flush the internal buffer by making appropriate API calls."""
+    async def _flush_buffer(self, chunks: Optional[Sequence[Union[Dict, Chunk]]] = None, **kwargs) -> AsyncSlackResponse:
+        """Flush the internal buffer with chunks by making appropriate API calls."""
+        flushings: List[Union[Dict, Chunk]] = []
+        if len(self._buffer) != 0:
+            flushings.append(MarkdownTextChunk(text=self._buffer))
+        if chunks is not None:
+            flushings.extend(chunks)
         if not self._stream_ts:
             response = await self._client.chat_startStream(
                 **self._stream_args,
                 token=self._token,
                 **kwargs,
-                markdown_text=self._buffer,
+                chunks=flushings,
             )
             self._stream_ts = response.get("ts")
             self._state = "in_progress"
@@ -206,7 +222,7 @@ async def _flush_buffer(self, **kwargs) -> AsyncSlackResponse:
                 channel=self._stream_args["channel"],
                 ts=self._stream_ts,
                 **kwargs,
-                markdown_text=self._buffer,
+                chunks=flushings,
             )
         self._buffer = ""
         return response

slack_sdk/web/async_client.py

@@ -17,12 +17,13 @@
 from typing import Any, Dict, List, Optional, Sequence, Union
 
 import slack_sdk.errors as e
+from slack_sdk.models.messages.chunk import Chunk
 from slack_sdk.models.views import View
 from slack_sdk.web.async_chat_stream import AsyncChatStream
 
 from ..models.attachments import Attachment
 from ..models.blocks import Block, RichTextBlock
-from ..models.metadata import Metadata, EntityMetadata, EventAndEntityMetadata
+from ..models.metadata import EntityMetadata, EventAndEntityMetadata, Metadata
 from .async_base_client import AsyncBaseClient, AsyncSlackResponse
 from .internal_utils import (
     _parse_web_class_objects,
@@ -2630,7 +2631,8 @@ async def chat_appendStream(
         *,
         channel: str,
         ts: str,
-        markdown_text: str,
+        markdown_text: Optional[str] = None,
+        chunks: Optional[Sequence[Union[Dict, Chunk]]] = None,
         **kwargs,
     ) -> AsyncSlackResponse:
         """Appends text to an existing streaming conversation.
@@ -2641,8 +2643,10 @@ async def chat_appendStream(
                 "channel": channel,
                 "ts": ts,
                 "markdown_text": markdown_text,
+                "chunks": chunks,
             }
         )
+        _parse_web_class_objects(kwargs)
         kwargs = _remove_none_values(kwargs)
         return await self.api_call("chat.appendStream", json=kwargs)
 
@@ -2884,6 +2888,7 @@ async def chat_startStream(
         markdown_text: Optional[str] = None,
         recipient_team_id: Optional[str] = None,
         recipient_user_id: Optional[str] = None,
+        chunks: Optional[Sequence[Union[Dict, Chunk]]] = None,
         **kwargs,
     ) -> AsyncSlackResponse:
         """Starts a new streaming conversation.
@@ -2896,8 +2901,10 @@ async def chat_startStream(
                 "markdown_text": markdown_text,
                 "recipient_team_id": recipient_team_id,
                 "recipient_user_id": recipient_user_id,
+                "chunks": chunks,
             }
         )
+        _parse_web_class_objects(kwargs)
         kwargs = _remove_none_values(kwargs)
         return await self.api_call("chat.startStream", json=kwargs)
 
@@ -2909,6 +2916,7 @@ async def chat_stopStream(
         markdown_text: Optional[str] = None,
         blocks: Optional[Union[str, Sequence[Union[Dict, Block]]]] = None,
         metadata: Optional[Union[Dict, Metadata]] = None,
+        chunks: Optional[Sequence[Union[Dict, Chunk]]] = None,
         **kwargs,
     ) -> AsyncSlackResponse:
         """Stops a streaming conversation.
@@ -2921,6 +2929,7 @@ async def chat_stopStream(
                 "markdown_text": markdown_text,
                 "blocks": blocks,
                 "metadata": metadata,
+                "chunks": chunks,
             }
         )
         _parse_web_class_objects(kwargs)