Skip to content

Commit

Permalink
feat: improve streaming (#67)
Browse files Browse the repository at this point in the history
* custom spinner

* tweaks

* remove ora

* refactor

* refactor
  • Loading branch information
mdjastrzebski authored Dec 7, 2024
1 parent 823f789 commit 2a3076f
Show file tree
Hide file tree
Showing 6 changed files with 102 additions and 136 deletions.
2 changes: 1 addition & 1 deletion package.json
Original file line number Diff line number Diff line change
Expand Up @@ -49,9 +49,9 @@
"@clack/prompts": "^0.8.2",
"ai": "^4.0.13",
"chalk": "^5.3.0",
"cli-spinners": "^3.2.0",
"date-fns": "^4.1.0",
"dotenv": "^16.4.7",
"ora": "^8.1.1",
"tiktoken": "^1.0.17",
"update-notifier": "^7.3.1",
"yargs": "^17.7.2",
Expand Down
21 changes: 0 additions & 21 deletions src/commands/chat/format.ts

This file was deleted.

35 changes: 25 additions & 10 deletions src/commands/chat/index.tsx
Original file line number Diff line number Diff line change
@@ -1,15 +1,15 @@
import { Application, createApp, Message } from '@callstack/byorg-core';
import { Application, AssistantResponse, createApp, Message } from '@callstack/byorg-core';
import type { CommandModule } from 'yargs';
import { checkIfConfigExists, parseConfigFile } from '../../config-file.js';
import { output, outputError, setVerbose } from '../../output.js';
import { getVerbose, output, outputError, setVerbose } from '../../output.js';
import { run as runInit } from '../init.js';
import { colorAssistant } from '../../colors.js';
import { colorAssistant, colorVerbose } from '../../colors.js';
import { formatSpeed, formatTime } from '../../format.js';
import { initInput, readUserInput, setInterruptHandler } from './input.js';
import { processChatCommand } from './commands.js';
import { cliOptions, type CliOptions } from './cli-options.js';
import { formatResponse } from './format.js';
import { getProvider, getProviderConfig, initProvider } from './providers.js';
import { spinnerStart, spinnerStop, spinnerUpdate } from './spinner.js';
import { streamingClear, streamingFinish, streamingStart, streamingUpdate } from './streaming.js';
import { messages } from './state.js';
import { texts } from './texts.js';
import { exit } from './utils.js';
Expand All @@ -34,6 +34,7 @@ async function run(initialPrompt: string, options: CliOptions) {
try {
const configFile = parseConfigFile();
initProvider(options, configFile);

const app = createApp({
chatModel: getProvider().getChatModel(getProviderConfig()),
systemPrompt: () => getProviderConfig().systemPrompt,
Expand All @@ -44,7 +45,7 @@ async function run(initialPrompt: string, options: CliOptions) {
}

setInterruptHandler(() => {
spinnerStop();
streamingClear();
exit();
});

Expand Down Expand Up @@ -73,19 +74,19 @@ async function processMessages(app: Application, messages: Message[]) {
const stream = getProviderConfig().stream;

const onPartialUpdate = (content: string) => {
spinnerUpdate(colorAssistant(`${texts.assistantLabel} ${content}`));
streamingUpdate(colorAssistant(`${texts.assistantLabel} ${content}`));
};

spinnerStart(colorAssistant(texts.assistantLabel));
streamingStart(colorAssistant(texts.assistantLabel));
const { response } = await app.processMessages(messages, {
onPartialResponse: stream ? onPartialUpdate : undefined,
});

if (response.role === 'assistant') {
messages.push({ role: 'assistant', content: response.content });
spinnerStop(`${formatResponse(response)}\n`);
streamingFinish(`${formatResponse(response)}\n`);
} else {
spinnerStop(response.content);
streamingFinish(response.content);
}

// Insert empty line after each response
Expand All @@ -99,3 +100,17 @@ function outputMessage(message: Message) {
output(colorAssistant(`${texts.assistantLabel} ${message.content}`));
}
}

export function formatResponse(response: AssistantResponse) {
let result = colorAssistant(`${texts.assistantLabel} ${response.content}`);

if (getVerbose()) {
const stats = `${formatTime(response.usage.responseTime)} ${formatSpeed(
response.usage?.outputTokens,
response.usage.responseTime,
)}`;
result += ` ${colorVerbose(stats)}`;
}

return result;
}
24 changes: 0 additions & 24 deletions src/commands/chat/spinner.ts

This file was deleted.

65 changes: 65 additions & 0 deletions src/commands/chat/streaming.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
import cliSpinners from 'cli-spinners';
import { colorAssistant } from '../../colors.js';

const spinner = cliSpinners.dots;
const frames = spinner.frames.map((f) => colorAssistant(f));

let currentLine = '';
let outputtedLines: string[] = [];

let intervalRef: NodeJS.Timeout | undefined;
let frameIndex = 0;

export function streamingStart(text: string) {
const lines = text.trimEnd().split('\n');
const linesToPrint = lines.slice(0, -1);
if (linesToPrint.length > 0) {
process.stdout.write(`${CLEAR_LINE}${linesToPrint.join('\n')}\n`);
}

outputtedLines = linesToPrint;
currentLine = lines[lines.length - 1];
startSpinner();
}

export function streamingUpdate(text: string) {
const lines = text.trimEnd().split('\n');
const linesToPrint = lines.slice(outputtedLines.length, -1);
if (linesToPrint.length > 0) {
process.stdout.write(`${CLEAR_LINE}${linesToPrint.join('\n')}\n`);
}

outputtedLines = lines.slice(0, -1);
currentLine = lines[lines.length - 1];
// Will render new text in the next animation frame
}

export function streamingFinish(text: string) {
clearInterval(intervalRef);
const lines = text.trimEnd().split('\n');
const linesToPrint = lines.slice(outputtedLines.length);
process.stdout.write(`${CLEAR_LINE}${linesToPrint.join('\n')}\n`);
}

export function streamingClear() {
clearInterval(intervalRef);
}

export const CLEAR_LINE = '\u001b[0G\u001b[2K';
export const DISABLE_WORD_WRAP = '\u001b[?7l';
export const ENABLE_WORD_WRAP = '\u001b[?7h';

export function startSpinner() {
if (intervalRef) {
clearInterval(intervalRef);
}

intervalRef = setInterval(renderFrame, spinner.interval).unref();
}

function renderFrame() {
frameIndex = (frameIndex + 1) % frames.length;
process.stdout.write(
`${CLEAR_LINE}${DISABLE_WORD_WRAP}${currentLine} ${frames[frameIndex]}${ENABLE_WORD_WRAP}`,
);
}
91 changes: 11 additions & 80 deletions yarn.lock
Original file line number Diff line number Diff line change
Expand Up @@ -558,6 +558,7 @@ __metadata:
"@vitest/coverage-v8": "npm:^2.1.8"
ai: "npm:^4.0.13"
chalk: "npm:^5.3.0"
cli-spinners: "npm:^3.2.0"
date-fns: "npm:^4.1.0"
del-cli: "npm:^6.0.0"
dotenv: "npm:^16.4.7"
Expand All @@ -569,7 +570,6 @@ __metadata:
jest: "npm:^29.7.0"
memfs: "npm:^4.14.1"
mock-fs: "npm:^5.4.1"
ora: "npm:^8.1.1"
prettier: "npm:^3.4.2"
release-it: "npm:^15.11.0"
tiktoken: "npm:^1.0.17"
Expand Down Expand Up @@ -3254,22 +3254,20 @@ __metadata:
languageName: node
linkType: hard

"cli-cursor@npm:^5.0.0":
version: 5.0.0
resolution: "cli-cursor@npm:5.0.0"
dependencies:
restore-cursor: "npm:^5.0.0"
checksum: 10/1eb9a3f878b31addfe8d82c6d915ec2330cec8447ab1f117f4aa34f0137fbb3137ec3466e1c9a65bcb7557f6e486d343f2da57f253a2f668d691372dfa15c090
languageName: node
linkType: hard

"cli-spinners@npm:^2.5.0, cli-spinners@npm:^2.6.1, cli-spinners@npm:^2.9.2":
"cli-spinners@npm:^2.5.0, cli-spinners@npm:^2.6.1":
version: 2.9.2
resolution: "cli-spinners@npm:2.9.2"
checksum: 10/a0a863f442df35ed7294424f5491fa1756bd8d2e4ff0c8736531d886cec0ece4d85e8663b77a5afaf1d296e3cbbebff92e2e99f52bbea89b667cbe789b994794
languageName: node
linkType: hard

"cli-spinners@npm:^3.2.0":
version: 3.2.0
resolution: "cli-spinners@npm:3.2.0"
checksum: 10/6612d3880c87ad1749556ff463c41499ebeab4024ee4afc41a8731d0bcd1679b18bb67a98df7e647cfa49adcff1ce86c049e141a4da028bb12831d7f13111d89
languageName: node
linkType: hard

"cli-width@npm:^4.0.0":
version: 4.1.0
resolution: "cli-width@npm:4.1.0"
Expand Down Expand Up @@ -6087,20 +6085,13 @@ __metadata:
languageName: node
linkType: hard

"is-unicode-supported@npm:^1.1.0, is-unicode-supported@npm:^1.2.0, is-unicode-supported@npm:^1.3.0":
"is-unicode-supported@npm:^1.1.0, is-unicode-supported@npm:^1.2.0":
version: 1.3.0
resolution: "is-unicode-supported@npm:1.3.0"
checksum: 10/20a1fc161afafaf49243551a5ac33b6c4cf0bbcce369fcd8f2951fbdd000c30698ce320de3ee6830497310a8f41880f8066d440aa3eb0a853e2aa4836dd89abc
languageName: node
linkType: hard

"is-unicode-supported@npm:^2.0.0":
version: 2.1.0
resolution: "is-unicode-supported@npm:2.1.0"
checksum: 10/f254e3da6b0ab1a57a94f7273a7798dd35d1d45b227759f600d0fa9d5649f9c07fa8d3c8a6360b0e376adf916d151ec24fc9a50c5295c58bae7ca54a76a063f9
languageName: node
linkType: hard

"is-weakmap@npm:^2.0.2":
version: 2.0.2
resolution: "is-weakmap@npm:2.0.2"
Expand Down Expand Up @@ -7023,16 +7014,6 @@ __metadata:
languageName: node
linkType: hard

"log-symbols@npm:^6.0.0":
version: 6.0.0
resolution: "log-symbols@npm:6.0.0"
dependencies:
chalk: "npm:^5.3.0"
is-unicode-supported: "npm:^1.3.0"
checksum: 10/510cdda36700cbcd87a2a691ea08d310a6c6b449084018f7f2ec4f732ca5e51b301ff1327aadd96f53c08318e616276c65f7fe22f2a16704fb0715d788bc3c33
languageName: node
linkType: hard

"loose-envify@npm:^1.4.0":
version: 1.4.0
resolution: "loose-envify@npm:1.4.0"
Expand Down Expand Up @@ -7227,13 +7208,6 @@ __metadata:
languageName: node
linkType: hard

"mimic-function@npm:^5.0.0":
version: 5.0.1
resolution: "mimic-function@npm:5.0.1"
checksum: 10/eb5893c99e902ccebbc267c6c6b83092966af84682957f79313311edb95e8bb5f39fb048d77132b700474d1c86d90ccc211e99bae0935447a4834eb4c882982c
languageName: node
linkType: hard

"mimic-response@npm:^3.1.0":
version: 3.1.0
resolution: "mimic-response@npm:3.1.0"
Expand Down Expand Up @@ -7660,15 +7634,6 @@ __metadata:
languageName: node
linkType: hard

"onetime@npm:^7.0.0":
version: 7.0.0
resolution: "onetime@npm:7.0.0"
dependencies:
mimic-function: "npm:^5.0.0"
checksum: 10/eb08d2da9339819e2f9d52cab9caf2557d80e9af8c7d1ae86e1a0fef027d00a88e9f5bd67494d350df360f7c559fbb44e800b32f310fb989c860214eacbb561c
languageName: node
linkType: hard

"open@npm:9.1.0":
version: 9.1.0
resolution: "open@npm:9.1.0"
Expand Down Expand Up @@ -7743,23 +7708,6 @@ __metadata:
languageName: node
linkType: hard

"ora@npm:^8.1.1":
version: 8.1.1
resolution: "ora@npm:8.1.1"
dependencies:
chalk: "npm:^5.3.0"
cli-cursor: "npm:^5.0.0"
cli-spinners: "npm:^2.9.2"
is-interactive: "npm:^2.0.0"
is-unicode-supported: "npm:^2.0.0"
log-symbols: "npm:^6.0.0"
stdin-discarder: "npm:^0.2.2"
string-width: "npm:^7.2.0"
strip-ansi: "npm:^7.1.0"
checksum: 10/2308c0a4da83099b0778a914890f5561329e3c948f73e865e0dced9053d62327219d8bde0f9b7c79a02913a7e41458ff5667d0115032fedc54907862e2de9695
languageName: node
linkType: hard

"os-name@npm:5.1.0":
version: 5.1.0
resolution: "os-name@npm:5.1.0"
Expand Down Expand Up @@ -8532,16 +8480,6 @@ __metadata:
languageName: node
linkType: hard

"restore-cursor@npm:^5.0.0":
version: 5.1.0
resolution: "restore-cursor@npm:5.1.0"
dependencies:
onetime: "npm:^7.0.0"
signal-exit: "npm:^4.1.0"
checksum: 10/838dd54e458d89cfbc1a923b343c1b0f170a04100b4ce1733e97531842d7b440463967e521216e8ab6c6f8e89df877acc7b7f4c18ec76e99fb9bf5a60d358d2c
languageName: node
linkType: hard

"retry@npm:0.13.1":
version: 0.13.1
resolution: "retry@npm:0.13.1"
Expand Down Expand Up @@ -8874,7 +8812,7 @@ __metadata:
languageName: node
linkType: hard

"signal-exit@npm:^4.0.1, signal-exit@npm:^4.1.0":
"signal-exit@npm:^4.0.1":
version: 4.1.0
resolution: "signal-exit@npm:4.1.0"
checksum: 10/c9fa63bbbd7431066174a48ba2dd9986dfd930c3a8b59de9c29d7b6854ec1c12a80d15310869ea5166d413b99f041bfa3dd80a7947bcd44ea8e6eb3ffeabfa1f
Expand Down Expand Up @@ -9057,13 +8995,6 @@ __metadata:
languageName: node
linkType: hard

"stdin-discarder@npm:^0.2.2":
version: 0.2.2
resolution: "stdin-discarder@npm:0.2.2"
checksum: 10/642ffd05bd5b100819d6b24a613d83c6e3857c6de74eb02fc51506fa61dc1b0034665163831873868157c4538d71e31762bcf319be86cea04c3aba5336470478
languageName: node
linkType: hard

"stop-iteration-iterator@npm:^1.0.0":
version: 1.0.0
resolution: "stop-iteration-iterator@npm:1.0.0"
Expand Down

0 comments on commit 2a3076f

Please sign in to comment.