From 91a63f2ccd6468a2b32d66f4ee657ac1a3acf4c6 Mon Sep 17 00:00:00 2001
From: Mike Amy <mikeamycoder@gmail.com>
Date: Tue, 9 May 2023 01:27:54 +0700
Subject: [PATCH] Fixed CLI streamed chat completions. (#319)

* Fixed streamed chat completions.

Streamed chat completions use a different response structure per returned token, also they may have roles and empty tokens at the end. Handle this sensibly.

* Only render content

---------

Co-authored-by: Atty Eleti <atty@openai.com>
---
 openai/cli.py | 24 ++++++++++++++++++------
 1 file changed, 18 insertions(+), 6 deletions(-)

diff --git a/openai/cli.py b/openai/cli.py
index e924133d72..ad08ac3e7b 100644
--- a/openai/cli.py
+++ b/openai/cli.py
@@ -141,9 +141,14 @@ def create(cls, args):
             for c_idx, c in enumerate(sorted(choices, key=lambda s: s["index"])):
                 if len(choices) > 1:
                     sys.stdout.write("===== Chat Completion {} =====\n".format(c_idx))
-                sys.stdout.write(c["message"]["content"])
-                if len(choices) > 1:
-                    sys.stdout.write("\n")
+                if args.stream:
+                    delta = c["delta"]
+                    if "content" in delta:
+                        sys.stdout.write(delta["content"])
+                else:
+                    sys.stdout.write(c["message"]["content"])
+                    if len(choices) > 1:  # not in streams
+                        sys.stdout.write("\n")
                 sys.stdout.flush()
 
 
@@ -203,7 +208,9 @@ def list(cls, args):
 
     @classmethod
     def create(cls, args):
-        models = openai.Deployment.create(model=args.model, scale_settings={"scale_type": args.scale_type})
+        models = openai.Deployment.create(
+            model=args.model, scale_settings={"scale_type": args.scale_type}
+        )
         print(models)
 
 
@@ -833,10 +840,15 @@ def help(args):
     sub = subparsers.add_parser("deployments.delete")
     sub.add_argument("-i", "--id", required=True, help="The deployment ID")
     sub.set_defaults(func=Deployment.delete)
-    
+
     sub = subparsers.add_parser("deployments.create")
     sub.add_argument("-m", "--model", required=True, help="The model ID")
-    sub.add_argument("-s", "--scale_type", required=True, help="The scale type. Either 'manual' or 'standard'")
+    sub.add_argument(
+        "-s",
+        "--scale_type",
+        required=True,
+        help="The scale type. Either 'manual' or 'standard'",
+    )
     sub.set_defaults(func=Deployment.create)
 
     # Models