Add option to disable response streaming (TheR1D#290)

Author: danarth

Diff for: README.md

@@ -313,6 +313,8 @@ DEFAULT_COLOR=magenta
 SYSTEM_ROLES=false
 # When in --shell mode, default to "Y" for no input.
 DEFAULT_EXECUTE_SHELL_CMD=false
+# Disable streaming of responses
+DISABLE_STREAMING=false
 ```
 Possible options for `DEFAULT_COLOR`: black, red, green, yellow, blue, magenta, cyan, white, bright_black, bright_red, bright_green, bright_yellow, bright_blue, bright_magenta, bright_cyan, bright_white.
 

Diff for: sgpt/client.py

@@ -10,6 +10,7 @@
 CACHE_LENGTH = int(cfg.get("CACHE_LENGTH"))
 CACHE_PATH = Path(cfg.get("CACHE_PATH"))
 REQUEST_TIMEOUT = int(cfg.get("REQUEST_TIMEOUT"))
+DISABLE_STREAMING = str(cfg.get("DISABLE_STREAMING"))
 
 
 class OpenAIClient:
@@ -37,12 +38,13 @@ def _request(
         :param top_probability: Float in 0.0 - 1.0 range.
         :return: Response body JSON.
         """
+        stream = DISABLE_STREAMING == "false"
         data = {
             "messages": messages,
             "model": model,
             "temperature": temperature,
             "top_p": top_probability,
-            "stream": True,
+            "stream": stream,
         }
         endpoint = f"{self.api_host}/v1/chat/completions"
         response = requests.post(
@@ -54,11 +56,15 @@ def _request(
             },
             json=data,
             timeout=REQUEST_TIMEOUT,
-            stream=True,
+            stream=stream,
         )
         response.raise_for_status()
         # TODO: Optimise.
         # https://github.com/openai/openai-python/blob/237448dc072a2c062698da3f9f512fae38300c1c/openai/api_requestor.py#L98
+        if not stream:
+            data = response.json()
+            yield data["choices"][0]["message"]["content"]  # type: ignore
+            return
         for line in response.iter_lines():
             data = line.lstrip(b"data: ").decode("utf-8")
             if data == "[DONE]":  # type: ignore

Diff for: sgpt/config.py

@@ -29,6 +29,7 @@
     "ROLE_STORAGE_PATH": os.getenv("ROLE_STORAGE_PATH", str(ROLE_STORAGE_PATH)),
     "SYSTEM_ROLES": os.getenv("SYSTEM_ROLES", "false"),
     "DEFAULT_EXECUTE_SHELL_CMD": os.getenv("DEFAULT_EXECUTE_SHELL_CMD", "false"),
+    "DISABLE_STREAMING": os.getenv("DISABLE_STREAMING", "false")
     # New features might add their own config variables here.
 }
 

Diff for: sgpt/handlers/handler.py

@@ -27,8 +27,11 @@ def get_completion(self, **kwargs: Any) -> Generator[str, None, None]:
     def handle(self, prompt: str, **kwargs: Any) -> str:
         messages = self.make_messages(self.make_prompt(prompt))
         full_completion = ""
+        stream = cfg.get("DISABLE_STREAMING") == "false"
+        if not stream:
+            typer.echo("Loading...\r", nl=False)
         for word in self.get_completion(messages=messages, **kwargs):
             typer.secho(word, fg=self.color, bold=True, nl=False)
             full_completion += word
-        typer.echo()
+        typer.echo("\033[K" if not stream else "")  # Overwrite "loading..."
         return full_completion

Diff for: tests/test_integration.py

@@ -30,6 +30,11 @@
 
 
 class TestShellGpt(TestCase):
+    @classmethod
+    def setUpClass(cls):
+        # Response streaming should be enabled for these tests.
+        assert cfg.get("DISABLE_STREAMING") == "false"
+
     def setUp(self) -> None:
         # Just to not spam the API.
         sleep(1)