add meta to CompletionChunkEvent

griptape-ai · Oct 17, 2024 · 1ca46e2 · 1ca46e2
1 parent db492c9
commit 1ca46e2
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/griptape/drivers/prompt/base_prompt_driver.py b/griptape/drivers/prompt/base_prompt_driver.py
@@ -127,12 +127,13 @@ def __process_stream(self, prompt_stack: PromptStack) -> Message:
                 else:
                     delta_contents[content.index] = [content]
                 if isinstance(content, TextDeltaMessageContent):
-                    EventBus.publish_event(CompletionChunkEvent(token=content.text))
+                    EventBus.publish_event(CompletionChunkEvent(token=content.text, meta={"type": "text"}))
                 elif isinstance(content, ActionCallDeltaMessageContent):
+                    meta = {"type": "action"}
                     if content.tag is not None and content.name is not None and content.path is not None:
-                        EventBus.publish_event(CompletionChunkEvent(token=str(content)))
+                        EventBus.publish_event(CompletionChunkEvent(token=str(content), meta=meta))
                     elif content.partial_input is not None:
-                        EventBus.publish_event(CompletionChunkEvent(token=content.partial_input))
+                        EventBus.publish_event(CompletionChunkEvent(token=content.partial_input, meta=meta))
 
         # Build a complete content from the content deltas
         return self.__build_message(list(delta_contents.values()), usage)