feat: Added use streaming field

writer · Nov 15, 2024 · 0e05573 · 0e05573
1 parent 10a7c5d
commit 0e05573
Showing 1 changed file with 20 additions and 5 deletions.
diff --git a/src/writer/workflows_blocks/writerchat.py b/src/writer/workflows_blocks/writerchat.py
@@ -21,6 +21,15 @@ def register(cls, type: str):
                         "desc": "Where the conversation will be stored",
                         "type": "Text",
                     },
+                    "useStreaming": {
+                        "name": "Use streaming",
+                        "type": "Text",
+                        "default": "yes",
+                        "options": {
+                            "yes": "Yes",
+                            "no": "No"
+                        }
+                    },
                     "tools": {
                         "name": "Tools",
                         "type": "Tools",
@@ -78,6 +87,7 @@ def run(self):
             import writer.ai
 
             conversation_state_element = self._get_field("conversationStateElement")
+            use_streaming = self._get_field("useStreaming", False, "yes") == "yes"
             tools_raw = self._get_field("tools", True)
             tools = []
 
@@ -110,12 +120,17 @@ def run(self):
 
             try:
                 msg = ""
-                for chunk in conversation.stream_complete(tools=tools):
-                    if chunk.get("content") is None:
-                        chunk["content"] = ""
-                    msg += chunk.get("content")
-                    conversation += chunk
+                if not use_streaming:
+                    msg = conversation.complete(tools=tools)
+                    conversation += msg
                     self.evaluator.set_state(conversation_state_element, self.instance_path, conversation, base_context=self.execution_env)
+                else:
+                    for chunk in conversation.stream_complete(tools=tools):
+                        if chunk.get("content") is None:
+                            chunk["content"] = ""
+                        msg += chunk.get("content")
+                        conversation += chunk
+                        self.evaluator.set_state(conversation_state_element, self.instance_path, conversation, base_context=self.execution_env)
             except BaseException:
                 msg = {
                     "role": "assistant",