process initial message

callstack · Dec 6, 2024 · 2dee726 · 2dee726
1 parent 4c96767
commit 2dee726
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 34 deletions.
diff --git a/src/commands/chat/index.tsx b/src/commands/chat/index.tsx
@@ -1,11 +1,13 @@
 import type { CommandModule } from 'yargs';
+import { Message } from '@callstack/byorg-core';
 import { parseConfigFile } from '../../config-file.js';
 import {
   clearCurrentLine,
   output,
   outputAssistant,
   outputAssistantPartial,
   outputError,
+  outputVerbose,
   readUserInput,
   setInterruptHandler,
   setVerbose,
@@ -41,8 +43,12 @@ async function run(initialPrompt: string, options: PromptOptions) {
       shouldExit = true;
     });
 
-    const { app, messages, stream, verbose } = createState(options, configFile, initialPrompt);
-    setVerbose(verbose);
+    const { app, messages, stream } = createState(options, configFile, initialPrompt);
+
+    if (messages.length > 0) {
+      messages.forEach((m) => outputMessage(m));
+      await processMessages(messages, { app, stream });
+    }
 
     while (!shouldExit) {
       const userMessage = await readUserInput();
@@ -51,35 +57,51 @@ async function run(initialPrompt: string, options: PromptOptions) {
       }
 
       messages.push({ role: 'user', content: userMessage });
+      await processMessages(messages, { app, stream });
+    }
+  } catch (error) {
+    outputError(error);
+    process.exit(1);
+  }
+}
 
-      lastMessage = '';
-      if (stream) {
-        outputAssistantPartial('ai: ');
-      }
+type ProcessMessagesOptions = {
+  app: App;
+  stream: boolean;
+};
 
-      const { response } = await app.processMessages(messages, {
-        onPartialResponse: stream ? onPartialUpdate : undefined,
-      });
+async function processMessages(messages: Message[], { app, stream }: ProcessMessagesOptions) {
+  lastMessage = '';
+  if (stream) {
+    outputAssistantPartial('ai: ');
+  }
 
-      if (stream) {
-        outputAssistantPartial('\n');
-      }
+  const { response } = await app.processMessages(messages, {
+    onPartialResponse: stream ? onPartialUpdate : undefined,
+  });
 
-      // With streaming, assistant response is already printed
-      if (response.role === 'assistant') {
-        messages.push({ role: 'assistant', content: response.content });
+  if (stream) {
+    outputAssistantPartial('\n');
+  }
 
-        if (!stream) {
-          outputAssistant(response.content);
-        }
-      }
+  // With streaming, assistant response is already printed
+  if (response.role === 'assistant') {
+    messages.push({ role: 'assistant', content: response.content });
 
-      if (response.role === 'system') {
-        output(response.content);
-      }
+    if (!stream) {
+      outputAssistant(response.content);
     }
-  } catch (error) {
-    outputError(error);
-    process.exit(1);
+  }
+
+  if (response.role === 'system') {
+    output(response.content);
+  }
+}
+
+function outputMessage(message: Message) {
+  if (message.role === 'user') {
+    output(`user: ${message.content}`);
+  } else {
+    outputAssistant(`ai: ${message.content}`);
   }
 }
diff --git a/src/commands/chat/state.ts b/src/commands/chat/state.ts
@@ -2,10 +2,10 @@ import { createApp, type Message } from '@callstack/byorg-core';
 import { type ConfigFile } from '../../config-file.js';
 import { DEFAULT_SYSTEM_PROMPT } from '../../default-config.js';
 import type { ResponseStyle } from '../../engine/providers/config.js';
+import { output, outputVerbose, outputWarning } from '../../output/index.js';
 import type { PromptOptions } from './prompt-options.js';
 import { getDefaultProvider, resolveProviderFromOption } from './providers.js';
 import { filterOutApiKey, handleInputFile } from './utils.js';
-import { output, outputVerbose, outputWarning } from '../../output/index.js';
 
 export let messages: Message[] = [];
 
@@ -56,6 +56,8 @@ export function createState(options: PromptOptions, configFile: ConfigFile, init
     }
   }
 
+  output('');
+
   if (initialPrompt) {
     messages.push({ role: 'user', content: initialPrompt });
   }

diff --git a/src/output/index.ts b/src/output/index.ts
@@ -45,21 +45,21 @@ export function outputError(error: unknown, ...args: unknown[]) {
   console.error(chalk.red(`ERROR: ${message}`, ...args));
 }
 
-export function outputAssistant(text: string) {
-  console.log(chalk.hex(colors.assistant)(text));
+export function outputAssistant(text: string, ...args: unknown[]) {
+  console.log(chalk.hex(colors.assistant)(text, ...args));
 }
 
-export function outputAssistantPartial(text: string) {
-  process.stdout.write(chalk.hex(colors.assistant)(text));
+export function outputAssistantPartial(text: string, ...args: unknown[]) {
+  process.stdout.write(chalk.hex(colors.assistant)(text, ...args));
 }
 
-export function output(text: string) {
-  console.log(chalk.hex(colors.info)(text));
+export function output(text: string, ...args: unknown[]) {
+  console.log(text, ...args);
 }
 
-export function outputVerbose(text: string) {
+export function outputVerbose(message: string, ...args: unknown[]) {
   if (showVerbose) {
-    console.log(chalk.dim(text));
+    console.log(chalk.dim(message, ...args));
   }
 }