fix: restore intro message and costs calculation (#69)

* fix: restore intro message and costs calculation

* debug command

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);
 }

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,37 +56,41 @@ 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() {
   const provider = getProvider();
   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);
 }

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);
   }

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);
     }
   }
 }

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;
+}

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;

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<ModelUsage>, 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<ModelUsage>, { provider, providerConfig }: CostContext) {
+  return (
+    provider.modelPricing[usage.model ?? providerConfig.model] ??
+    provider.modelPricing[providerConfig.model]
+  );
+}

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;

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`;

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));