Skip to content

Commit 70327ca

Browse files
zimegmwbrooks
andauthored
feat: add plan_update chunk for chat streaming (#1821)
Co-authored-by: Michael Brooks <[email protected]>
1 parent b1dc199 commit 70327ca

File tree

3 files changed

+49
-5
lines changed

3 files changed

+49
-5
lines changed

slack_sdk/models/messages/chunk.py

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@ def parse(cls, chunk: Union[Dict, "Chunk"]) -> Optional["Chunk"]:
3434
type = chunk["type"]
3535
if type == MarkdownTextChunk.type:
3636
return MarkdownTextChunk(**chunk)
37+
elif type == PlanUpdateChunk.type:
38+
return PlanUpdateChunk(**chunk)
3739
elif type == TaskUpdateChunk.type:
3840
return TaskUpdateChunk(**chunk)
3941
else:
@@ -67,6 +69,29 @@ def __init__(
6769
self.text = text
6870

6971

72+
class PlanUpdateChunk(Chunk):
73+
type = "plan_update"
74+
75+
@property
76+
def attributes(self) -> Set[str]: # type: ignore[override]
77+
return super().attributes.union({"title"})
78+
79+
def __init__(
80+
self,
81+
*,
82+
title: str,
83+
**others: Dict,
84+
):
85+
"""An updated title of plans for task and tool calls.
86+
87+
https://docs.slack.dev/messaging/sending-and-scheduling-messages#text-streaming
88+
"""
89+
super().__init__(type=self.type)
90+
show_unknown_key_warning(self, others)
91+
92+
self.title = title
93+
94+
7095
class TaskUpdateChunk(Chunk):
7196
type = "task_update"
7297

tests/slack_sdk/models/test_chunks.py

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import unittest
22

33
from slack_sdk.models.blocks.block_elements import UrlSourceElement
4-
from slack_sdk.models.messages.chunk import MarkdownTextChunk, TaskUpdateChunk
4+
from slack_sdk.models.messages.chunk import MarkdownTextChunk, PlanUpdateChunk, TaskUpdateChunk
55

66

77
class MarkdownTextChunkTests(unittest.TestCase):
@@ -15,6 +15,17 @@ def test_json(self):
1515
)
1616

1717

18+
class PlanUpdateChunkTests(unittest.TestCase):
19+
def test_json(self):
20+
self.assertDictEqual(
21+
PlanUpdateChunk(title="Crunching numbers...").to_dict(),
22+
{
23+
"type": "plan_update",
24+
"title": "Crunching numbers...",
25+
},
26+
)
27+
28+
1829
class TaskUpdateChunkTests(unittest.TestCase):
1930
def test_json(self):
2031
self.assertDictEqual(

tests/slack_sdk/web/test_internal_utils.py

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,9 @@
44
from pathlib import Path
55
from typing import Dict
66

7-
87
from slack_sdk.models.attachments import Attachment
98
from slack_sdk.models.blocks import Block, DividerBlock
10-
from slack_sdk.models.messages.chunk import MarkdownTextChunk, TaskUpdateChunk
9+
from slack_sdk.models.messages.chunk import MarkdownTextChunk, PlanUpdateChunk, TaskUpdateChunk
1110
from slack_sdk.web.internal_utils import (
1211
_build_unexpected_body_error_message,
1312
_get_url,
@@ -59,9 +58,14 @@ def test_can_parse_sequence_of_attachments(self):
5958

6059
def test_can_parse_sequence_of_chunks(self):
6160
for chunks in [
62-
[MarkdownTextChunk(text="fiz"), TaskUpdateChunk(id="001", title="baz", status="complete")], # list
61+
[
62+
MarkdownTextChunk(text="fiz"),
63+
PlanUpdateChunk(title="fuz"),
64+
TaskUpdateChunk(id="001", title="baz", status="complete"),
65+
], # list
6366
(
6467
MarkdownTextChunk(text="fiz"),
68+
PlanUpdateChunk(title="fuz"),
6569
TaskUpdateChunk(id="001", title="baz", status="complete"),
6670
), # tuple
6771
]:
@@ -87,7 +91,11 @@ def test_can_parse_str_attachments(self):
8791

8892
def test_can_parse_str_chunks(self):
8993
input = json.dumps(
90-
[MarkdownTextChunk(text="fiz").to_dict(), TaskUpdateChunk(id="001", title="baz", status="complete").to_dict()]
94+
[
95+
MarkdownTextChunk(text="fiz").to_dict(),
96+
PlanUpdateChunk(title="fuz").to_dict(),
97+
TaskUpdateChunk(id="001", title="baz", status="complete").to_dict(),
98+
]
9199
)
92100
kwargs = {"chunks": input}
93101
_parse_web_class_objects(kwargs)

0 commit comments

Comments
 (0)