Add streaming and JSON support to cli prompt command

ts/packages/aiclient/src/openai.ts

@@ -486,7 +486,6 @@ export function createChatModel(
         let completionParams: CompletionSettings | undefined;
         if (completionSettings) {
             completionParams = { ...completionSettings };
-            delete completionParams.response_format;
         }
         const params = {
             ...defaultParams,

ts/packages/cli/package.json

@@ -47,6 +47,7 @@
     "@oclif/plugin-help": "^5",
     "agent-cache": "workspace:*",
     "agent-dispatcher": "workspace:*",
+    "aiclient": "workspace:*",
     "chalk": "^5.3.0",
     "common-utils": "workspace:*",
     "debug": "^4.3.4",

ts/packages/cli/src/commands/prompt.ts

@@ -1,8 +1,9 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-import { Args, Command } from "@oclif/core";
-import { createLanguageModel } from "typechat";
+import { Args, Command, Flags } from "@oclif/core";
+import { getChatModelNames } from "common-utils";
+import { openai } from "aiclient";
 import fs from "node:fs";
 import chalk from "chalk";
 
@@ -14,10 +15,29 @@ export default class Prompt extends Command {
             required: true,
         }),
     };
+    static flags = {
+        model: Flags.string({
+            description: "Model to use",
+            options: getChatModelNames(),
+        }),
+        stream: Flags.boolean({
+            description: "Whether to stream the result",
+        }),
+        json: Flags.boolean({
+            description: "Output JSON response",
+        }),
+    };
 
     async run(): Promise<void> {
-        const { args } = await this.parse(Prompt);
-        const model = createLanguageModel(process.env);
+        const { args, flags } = await this.parse(Prompt);
+        const model = openai.createChatModel(
+            flags.model,
+            flags.json
+                ? {
+                      response_format: { type: "json_object" },
+                  }
+                : undefined,
+        );
         const isFile = fs.existsSync(args.request);
         if (isFile) {
             console.log(`Loading prompt from file: ${args.request}`);
@@ -31,18 +51,66 @@ export default class Prompt extends Command {
         console.log(`Sending prompt to GPT:`);
 
         console.log(chalk.grey(request));
-        const complete = await model.complete(request);
-        if (complete.success) {
-            console.log(chalk.green("GPT response:"));
+        console.log();
+        let responseText = "";
+        let time = 0;
+        try {
+            if (flags.stream) {
+                console.log(chalk.green("Streaming GPT response:"));
+                let number_chunks = 0;
+                const start = performance.now();
+                let first = start;
+                for await (const chunk of model.completeStream(request)) {
+                    if (first === start) {
+                        first = performance.now();
+                    }
+                    process.stdout.write(chalk.gray(chunk));
+                    responseText += chunk;
+                    number_chunks++;
+                }
+                const end = performance.now();
+                console.log();
+                console.log(
+                    `${number_chunks} chunks streamed in ${end - start}ms (first response in ${first - start}ms)`,
+                );
+                console.log();
+                time = end - start;
+            } else {
+                const start = performance.now();
+                const complete = await model.complete(request);
+                const end = performance.now();
+                if (complete.success) {
+                    responseText = complete.data;
+                    if (!flags.json) {
+                        console.log(
+                            chalk.green(
+                                `Full GPT response: (${end - start}ms)`,
+                            ),
+                        );
+                        console.log(responseText);
+                    }
+                    time = end - start;
+                } else {
+                    throw new Error(complete.message);
+                }
+            }
+        } catch (e: any) {
+            console.log("GPT error:");
+            console.log(e.message);
+            return;
+        }
+
+        if (flags.json) {
             try {
-                const json = JSON.parse(complete.data);
+                const json = JSON.parse(responseText);
+                console.log(chalk.green(`GPT JSON response: (${time}ms)`));
                 console.log(JSON.stringify(json, undefined, 2));
             } catch {
-                console.log(complete.data);
+                console.log(chalk.red("Failed to parse JSON response:"));
+                if (!flags.stream) {
+                    console.log(responseText);
+                }
             }
-        } else {
-            console.log("GPT error:");
-            console.log(complete.message);
         }
     }
 }