From 0e05573a8f035acffd60f75478946dbf5b677408 Mon Sep 17 00:00:00 2001
From: Ramiro Medina <64783088+ramedina86@users.noreply.github.com>
Date: Fri, 15 Nov 2024 14:27:42 +0000
Subject: [PATCH] feat: Added use streaming field

---
 src/writer/workflows_blocks/writerchat.py | 25 ++++++++++++++++++-----
 1 file changed, 20 insertions(+), 5 deletions(-)

diff --git a/src/writer/workflows_blocks/writerchat.py b/src/writer/workflows_blocks/writerchat.py
index b0c7814ea..02afe9558 100644
--- 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",