From f348057c4f2497affe3e8c23fc4a833493fc4aa8 Mon Sep 17 00:00:00 2001 From: Maciej Jastrzebski Date: Sun, 8 Dec 2024 09:50:35 +0100 Subject: [PATCH] fix: restore intro message and costs calculation (#69) * fix: restore intro message and costs calculation * debug command --- src/colors.ts | 4 +++ src/commands/chat/commands.ts | 62 ++++++++++++++++++++------------- src/commands/chat/index.tsx | 9 +++-- src/commands/chat/providers.tsx | 6 ++-- src/commands/chat/state.ts | 26 +++++++++++++- src/commands/chat/texts.ts | 3 +- src/commands/chat/usage.ts | 27 ++++++++++++++ src/commands/chat/utils.ts | 10 ++---- src/engine/session.ts | 42 ---------------------- src/format.ts | 23 +----------- src/output.ts | 6 +++- 11 files changed, 112 insertions(+), 106 deletions(-) create mode 100644 src/commands/chat/usage.ts delete mode 100644 src/engine/session.ts diff --git a/src/colors.ts b/src/colors.ts index 69858a9..ad26c56 100644 --- a/src/colors.ts +++ b/src/colors.ts @@ -17,6 +17,10 @@ export function colorWarning(...text: unknown[]) { return chalk.hex(colors.warning)(...text); } +export function colorSystem(...text: unknown[]) { + return chalk.grey(...text); +} + export function colorVerbose(...text: unknown[]) { return chalk.dim(...text); } diff --git a/src/commands/chat/commands.ts b/src/commands/chat/commands.ts index f1cf7c1..5978d90 100644 --- a/src/commands/chat/commands.ts +++ b/src/commands/chat/commands.ts @@ -1,8 +1,10 @@ import { CHATS_SAVE_DIRECTORY } from '../../file-utils.js'; -import { getVerbose, output, outputVerbose, outputWarning, setVerbose } from '../../output.js'; +import { formatCost, formatSpeed, formatTokenCount } from '../../format.js'; +import { getVerbose, outputSystem, outputWarning, setVerbose } from '../../output.js'; import { getProvider, getProviderConfig } from './providers.js'; -import { messages } from './state.js'; -import { exit, saveConversation } from './utils.js'; +import { messages, totalUsage } from './state.js'; +import { calculateUsageCost } from './usage.js'; +import { exit, filterOutApiKey, saveConversation } from './utils.js'; export function processChatCommand(input: string) { if (!input.startsWith('/')) { @@ -24,21 +26,27 @@ export function processChatCommand(input: string) { return true; } + if (command === '/debug') { + outputDebugInfo(); + return true; + } + if (command === '/forget') { // Clear all messages messages.length = 0; + outputSystem('Forgot all past messages from the current session.\n'); return true; } if (command === '/verbose') { setVerbose(!getVerbose()); - output(`Verbose mode: ${getVerbose() ? 'on' : 'off'}`); + outputSystem(`Verbose mode: ${getVerbose() ? 'on' : 'off'}\n`); return true; } if (input === '/save') { const saveConversationMessage = saveConversation(messages); - output(saveConversationMessage); + outputSystem(saveConversationMessage); return true; } @@ -48,17 +56,15 @@ export function processChatCommand(input: string) { export function outputHelp() { const lines = [ - '', 'Available commands:', - ' - /exit: Exit the CLI', - ' - /info: Show current provider, model, and system prompt', - ' - /forget: AI will forget previous messages', - ` - /save: Save in a text file in ${CHATS_SAVE_DIRECTORY}`, - ' - /verbose: Toggle verbose output', + '- /exit: Exit the CLI', + '- /info: Show current provider, model, and system prompt', + '- /forget: AI will forget previous messages', + `- /save: Save in a text file in ${CHATS_SAVE_DIRECTORY}`, + '- /verbose: Toggle verbose output', '', ]; - - output(lines.join('\n')); + outputSystem(lines.join('\n')); } export function outputInfo() { @@ -66,19 +72,25 @@ export function outputInfo() { const providerConfig = getProviderConfig(); const lines = [ - '', - 'Info:', - ` - Provider: ${provider.label}`, - ` - Model: ${providerConfig.model}`, - ` - System prompt: ${providerConfig.systemPrompt}`, + 'Session info:', + `- Provider: ${provider.label}`, + `- Model: ${providerConfig.model}`, + `- Cost: ${formatCost(calculateUsageCost(totalUsage, { provider, providerConfig }))}`, + `- Usage: ${formatTokenCount(totalUsage.inputTokens)} input token(s), ${formatTokenCount(totalUsage.outputTokens)} output token(s), ${totalUsage.requests} request(s)usag`, + `- Avg Speed: ${formatSpeed(totalUsage.outputTokens, totalUsage.responseTime)}`, + `- System prompt: ${providerConfig.systemPrompt}`, '', ]; - output(lines.join('\n')); + outputSystem(lines.join('\n')); +} + +export function outputDebugInfo() { + outputSystem(`Provider: ${toJson(getProvider().label)}\n`); + outputSystem(`Provider Config: ${toJson(getProviderConfig())}\n`); + outputSystem(`Messages: ${toJson(messages)}\n`); + outputSystem(`Usage: ${toJson(totalUsage)}\n`); +} - const rawMessages = JSON.stringify( - messages.map((m) => `${m.role}: ${m.content}`), - null, - 2, - ); - outputVerbose(`Messages: ${rawMessages}\n`); +function toJson(value: any) { + return JSON.stringify(value, filterOutApiKey, 2); } diff --git a/src/commands/chat/index.tsx b/src/commands/chat/index.tsx index 7d9ac5c..ce84460 100644 --- a/src/commands/chat/index.tsx +++ b/src/commands/chat/index.tsx @@ -1,7 +1,7 @@ import { Application, AssistantResponse, createApp, Message } from '@callstack/byorg-core'; import type { CommandModule } from 'yargs'; import { checkIfConfigExists, parseConfigFile } from '../../config-file.js'; -import { getVerbose, output, outputError, setVerbose } from '../../output.js'; +import { getVerbose, output, outputError, outputSystem, setVerbose } from '../../output.js'; import { run as runInit } from '../init.js'; import { colorAssistant, colorVerbose } from '../../colors.js'; import { formatSpeed, formatTime } from '../../format.js'; @@ -10,7 +10,7 @@ import { processChatCommand } from './commands.js'; import { cliOptions, type CliOptions } from './cli-options.js'; import { getProvider, getProviderConfig, initProvider } from './providers.js'; import { streamingClear, streamingFinish, streamingStart, streamingUpdate } from './streaming.js'; -import { messages } from './state.js'; +import { messages, updateUsage } from './state.js'; import { texts } from './texts.js'; import { exit } from './utils.js'; @@ -54,6 +54,8 @@ async function run(initialPrompt: string, options: CliOptions) { await processMessages(app, messages); } + outputSystem(texts.initialHelp); + // eslint-disable-next-line no-constant-condition while (true) { const userMessage = await readUserInput(); @@ -83,8 +85,9 @@ async function processMessages(app: Application, messages: Message[]) { }); if (response.role === 'assistant') { - messages.push({ role: 'assistant', content: response.content }); streamingFinish(`${formatResponse(response)}\n`); + messages.push({ role: 'assistant', content: response.content }); + updateUsage(response.usage); } else { streamingFinish(response.content); } diff --git a/src/commands/chat/providers.tsx b/src/commands/chat/providers.tsx index 6db7543..620b433 100644 --- a/src/commands/chat/providers.tsx +++ b/src/commands/chat/providers.tsx @@ -10,7 +10,7 @@ import openAi from '../../engine/providers/open-ai.js'; import anthropic from '../../engine/providers/anthropic.js'; import perplexity from '../../engine/providers/perplexity.js'; import mistral from '../../engine/providers/mistral.js'; -import { output, outputVerbose, outputWarning } from '../../output.js'; +import { outputSystem, outputVerbose, outputWarning } from '../../output.js'; import { CliOptions } from './cli-options.js'; import { filterOutApiKey, handleInputFile } from './utils.js'; @@ -97,13 +97,13 @@ export function initProvider(options: CliOptions, configFile: ConfigFile) { systemPrompt: fileSystemPrompt, costWarning, costInfo, - } = handleInputFile(options.file, getProviderConfig(), provider); + } = handleInputFile(options.file, providerConfig, provider); providerConfig.systemPrompt += `\n\n${fileSystemPrompt}`; if (costWarning) { outputWarning(costWarning); } else if (costInfo) { - output(costInfo); + outputSystem(costInfo); } } } diff --git a/src/commands/chat/state.ts b/src/commands/chat/state.ts index 0aebae9..8759b03 100644 --- a/src/commands/chat/state.ts +++ b/src/commands/chat/state.ts @@ -1,3 +1,27 @@ -import { type Message } from '@callstack/byorg-core'; +import { ModelUsage, type Message } from '@callstack/byorg-core'; export const messages: Message[] = []; + +export const totalUsage: ModelUsage = { + inputTokens: 0, + outputTokens: 0, + requests: 0, + responseTime: 0, + model: '', + usedTools: {}, +}; + +export function resetUsage() { + totalUsage.inputTokens = 0; + totalUsage.outputTokens = 0; + totalUsage.requests = 0; + totalUsage.responseTime = 0; +} + +export function updateUsage(usage: ModelUsage) { + totalUsage.inputTokens += usage.inputTokens; + totalUsage.outputTokens += usage.outputTokens; + totalUsage.requests += usage.requests; + totalUsage.responseTime += usage.responseTime; + totalUsage.model = usage.model; +} diff --git a/src/commands/chat/texts.ts b/src/commands/chat/texts.ts index 6726ac3..3b2080f 100644 --- a/src/commands/chat/texts.ts +++ b/src/commands/chat/texts.ts @@ -1,6 +1,5 @@ export const texts = { userLabel: 'me:', assistantLabel: 'ai:', - initialHelp: 'Type "/exit" or press Ctrl+C to exit. Type "/help" to see available commands.', - responseLoading: 'Thinking ...', + initialHelp: 'Type "/exit" or press Ctrl+C to exit. Type "/help" to see available commands.\n', } as const; diff --git a/src/commands/chat/usage.ts b/src/commands/chat/usage.ts new file mode 100644 index 0000000..4df19b4 --- /dev/null +++ b/src/commands/chat/usage.ts @@ -0,0 +1,27 @@ +import { ModelUsage } from '@callstack/byorg-core'; +import { Provider } from '../../engine/providers/provider.js'; +import { ProviderConfig } from '../../engine/providers/config.js'; + +export type CostContext = { + provider: Provider; + providerConfig: ProviderConfig; +}; + +export function calculateUsageCost(usage: Partial, context: CostContext) { + const pricing = getModelPricing(usage, context); + if (pricing === undefined) { + return undefined; + } + + const inputCost = ((usage.inputTokens ?? 0) * (pricing.inputTokensCost ?? 0)) / 1_000_000; + const outputCost = ((usage.outputTokens ?? 0) * (pricing.outputTokensCost ?? 0)) / 1_000_000; + const requestsCost = ((usage.requests ?? 0) * (pricing.requestsCost ?? 0)) / 1_000_000; + return inputCost + outputCost + requestsCost; +} + +function getModelPricing(usage: Partial, { provider, providerConfig }: CostContext) { + return ( + provider.modelPricing[usage.model ?? providerConfig.model] ?? + provider.modelPricing[providerConfig.model] + ); +} diff --git a/src/commands/chat/utils.ts b/src/commands/chat/utils.ts index 70f166c..546e37d 100644 --- a/src/commands/chat/utils.ts +++ b/src/commands/chat/utils.ts @@ -7,7 +7,6 @@ import { FILE_COST_WARNING, FILE_TOKEN_COUNT_WARNING, } from '../../default-config.js'; -import { calculateUsageCost } from '../../engine/session.js'; import { getTokensCount } from '../../engine/tokenizer.js'; import type { ProviderConfig } from '../../engine/providers/config.js'; import type { Provider } from '../../engine/providers/provider.js'; @@ -17,13 +16,12 @@ import { getDefaultFilename, getUniqueFilename, } from '../../file-utils.js'; -import { output } from '../../output.js'; import { texts } from './texts.js'; import { closeInput } from './input.js'; +import { calculateUsageCost } from './usage.js'; export function exit() { closeInput(); - output('\nBye...'); process.exit(0); } @@ -35,7 +33,7 @@ interface HandleInputFileResult { export function handleInputFile( inputFile: string, - config: ProviderConfig, + providerConfig: ProviderConfig, provider: Provider, ): HandleInputFileResult { const filePath = path.resolve(inputFile.replace('~', os.homedir())); @@ -46,9 +44,7 @@ export function handleInputFile( const fileContent = fs.readFileSync(filePath).toString(); const fileTokens = getTokensCount(fileContent); - - const pricing = provider.modelPricing[config.model]; - const fileCost = calculateUsageCost({ inputTokens: fileTokens }, pricing); + const fileCost = calculateUsageCost({ inputTokens: fileTokens }, { provider, providerConfig }); let costWarning = null; let costInfo = null; diff --git a/src/engine/session.ts b/src/engine/session.ts deleted file mode 100644 index 3b358bc..0000000 --- a/src/engine/session.ts +++ /dev/null @@ -1,42 +0,0 @@ -import type { ModelUsage } from '@callstack/byorg-core'; -import type { ModelPricing } from './providers/provider.js'; - -export interface SessionUsage { - total: ModelUsage; - current: ModelUsage; -} - -export interface SessionCost { - total: number; - current: number; -} - -export function combineUsage(usage1: ModelUsage, usage2: ModelUsage) { - return { - inputTokens: usage1.inputTokens + usage2.inputTokens, - outputTokens: usage1.outputTokens + usage2.outputTokens, - requests: usage1.requests + usage2.requests, - }; -} - -export function calculateUsageCost(usage: Partial, pricing: ModelPricing | undefined) { - if (pricing === undefined) { - return undefined; - } - - const inputCost = ((usage.inputTokens ?? 0) * (pricing.inputTokensCost ?? 0)) / 1_000_000; - const outputCost = ((usage.outputTokens ?? 0) * (pricing.outputTokensCost ?? 0)) / 1_000_000; - const requestsCost = ((usage.requests ?? 0) * (pricing.requestsCost ?? 0)) / 1_000_000; - return inputCost + outputCost + requestsCost; -} - -export function calculateSessionCost(usage: SessionUsage, pricing: ModelPricing | undefined) { - if (pricing === undefined) { - return undefined; - } - - return { - current: calculateUsageCost(usage.current, pricing)!, - total: calculateUsageCost(usage.total, pricing)!, - }; -} diff --git a/src/format.ts b/src/format.ts index e648e8f..0c758e3 100644 --- a/src/format.ts +++ b/src/format.ts @@ -1,5 +1,3 @@ -import type { SessionCost, SessionUsage } from './engine/session.js'; - export function formatCost(value: number | undefined, precision = 4) { if (value == null) { return '?'; @@ -29,32 +27,13 @@ export function formatTokenCount(tokenCount: number, roundTo = 1) { return `${scaledCount.toFixed(0)}${suffixes[suffixIndex]}`; } -export function formatSessionStats(responseTime?: number, usage?: SessionUsage) { - const parts = [ - responseTime ? `time: ${(responseTime / 1000).toFixed(1)} s` : undefined, - usage - ? `tokens: ${usage.current.inputTokens}+${usage.current.outputTokens} (total: ${usage.total.inputTokens}+${usage.total.outputTokens})` - : undefined, - ]; - - return parts.filter((x) => x !== undefined).join(', '); -} - -export function formatSessionCost(cost: SessionCost | undefined) { - if (cost === undefined) { - return undefined; - } - - return `costs: ${formatCost(cost.current)} (total: ${formatCost(cost.total)})`; -} - export function formatTime(timeInMs: number) { return `${(timeInMs / 1000).toFixed(1)} s`; } export function formatSpeed(tokens: number, timeInMs: number) { if (tokens == null || timeInMs == null || timeInMs === 0) { - return '? tok/s'; + return '? tokens/s'; } return `${((tokens * 1000) / timeInMs).toFixed(1)} tok/s`; diff --git a/src/output.ts b/src/output.ts index 4f004c7..d9b8b3a 100644 --- a/src/output.ts +++ b/src/output.ts @@ -1,4 +1,4 @@ -import { colorError, colorVerbose, colorWarning } from './colors.js'; +import { colorError, colorSystem, colorVerbose, colorWarning } from './colors.js'; let showVerbose = false; @@ -14,6 +14,10 @@ export function output(text: string, ...args: unknown[]) { console.log(text, ...args); } +export function outputSystem(text: string, ...args: unknown[]) { + console.log(colorSystem(text, ...args)); +} + export function outputVerbose(message: string, ...args: unknown[]) { if (showVerbose) { console.log(colorVerbose(message, ...args));