Skip to content

Commit f348057

Browse files
mdjastrzebskimdjastrzebski
authored
fix: restore intro message and costs calculation (#69)
* fix: restore intro message and costs calculation * debug command
1 parent d5c5de5 commit f348057

File tree

11 files changed

+112
-106
lines changed

11 files changed

+112
-106
lines changed

‎src/colors.ts

+4
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,10 @@ export function colorWarning(...text: unknown[]) {
1717
return chalk.hex(colors.warning)(...text);
1818
}
1919

20+
export function colorSystem(...text: unknown[]) {
21+
return chalk.grey(...text);
22+
}
23+
2024
export function colorVerbose(...text: unknown[]) {
2125
return chalk.dim(...text);
2226
}

‎src/commands/chat/commands.ts

+37-25
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
import { CHATS_SAVE_DIRECTORY } from '../../file-utils.js';
2-
import { getVerbose, output, outputVerbose, outputWarning, setVerbose } from '../../output.js';
2+
import { formatCost, formatSpeed, formatTokenCount } from '../../format.js';
3+
import { getVerbose, outputSystem, outputWarning, setVerbose } from '../../output.js';
34
import { getProvider, getProviderConfig } from './providers.js';
4-
import { messages } from './state.js';
5-
import { exit, saveConversation } from './utils.js';
5+
import { messages, totalUsage } from './state.js';
6+
import { calculateUsageCost } from './usage.js';
7+
import { exit, filterOutApiKey, saveConversation } from './utils.js';
68

79
export function processChatCommand(input: string) {
810
if (!input.startsWith('/')) {
@@ -24,21 +26,27 @@ export function processChatCommand(input: string) {
2426
return true;
2527
}
2628

29+
if (command === '/debug') {
30+
outputDebugInfo();
31+
return true;
32+
}
33+
2734
if (command === '/forget') {
2835
// Clear all messages
2936
messages.length = 0;
37+
outputSystem('Forgot all past messages from the current session.\n');
3038
return true;
3139
}
3240

3341
if (command === '/verbose') {
3442
setVerbose(!getVerbose());
35-
output(`Verbose mode: ${getVerbose() ? 'on' : 'off'}`);
43+
outputSystem(`Verbose mode: ${getVerbose() ? 'on' : 'off'}\n`);
3644
return true;
3745
}
3846

3947
if (input === '/save') {
4048
const saveConversationMessage = saveConversation(messages);
41-
output(saveConversationMessage);
49+
outputSystem(saveConversationMessage);
4250
return true;
4351
}
4452

@@ -48,37 +56,41 @@ export function processChatCommand(input: string) {
4856

4957
export function outputHelp() {
5058
const lines = [
51-
'',
5259
'Available commands:',
53-
' - /exit: Exit the CLI',
54-
' - /info: Show current provider, model, and system prompt',
55-
' - /forget: AI will forget previous messages',
56-
` - /save: Save in a text file in ${CHATS_SAVE_DIRECTORY}`,
57-
' - /verbose: Toggle verbose output',
60+
'- /exit: Exit the CLI',
61+
'- /info: Show current provider, model, and system prompt',
62+
'- /forget: AI will forget previous messages',
63+
`- /save: Save in a text file in ${CHATS_SAVE_DIRECTORY}`,
64+
'- /verbose: Toggle verbose output',
5865
'',
5966
];
60-
61-
output(lines.join('\n'));
67+
outputSystem(lines.join('\n'));
6268
}
6369

6470
export function outputInfo() {
6571
const provider = getProvider();
6672
const providerConfig = getProviderConfig();
6773

6874
const lines = [
69-
'',
70-
'Info:',
71-
` - Provider: ${provider.label}`,
72-
` - Model: ${providerConfig.model}`,
73-
` - System prompt: ${providerConfig.systemPrompt}`,
75+
'Session info:',
76+
`- Provider: ${provider.label}`,
77+
`- Model: ${providerConfig.model}`,
78+
`- Cost: ${formatCost(calculateUsageCost(totalUsage, { provider, providerConfig }))}`,
79+
`- Usage: ${formatTokenCount(totalUsage.inputTokens)} input token(s), ${formatTokenCount(totalUsage.outputTokens)} output token(s), ${totalUsage.requests} request(s)usag`,
80+
`- Avg Speed: ${formatSpeed(totalUsage.outputTokens, totalUsage.responseTime)}`,
81+
`- System prompt: ${providerConfig.systemPrompt}`,
7482
'',
7583
];
76-
output(lines.join('\n'));
84+
outputSystem(lines.join('\n'));
85+
}
86+
87+
export function outputDebugInfo() {
88+
outputSystem(`Provider: ${toJson(getProvider().label)}\n`);
89+
outputSystem(`Provider Config: ${toJson(getProviderConfig())}\n`);
90+
outputSystem(`Messages: ${toJson(messages)}\n`);
91+
outputSystem(`Usage: ${toJson(totalUsage)}\n`);
92+
}
7793

78-
const rawMessages = JSON.stringify(
79-
messages.map((m) => `${m.role}: ${m.content}`),
80-
null,
81-
2,
82-
);
83-
outputVerbose(`Messages: ${rawMessages}\n`);
94+
function toJson(value: any) {
95+
return JSON.stringify(value, filterOutApiKey, 2);
8496
}

‎src/commands/chat/index.tsx

+6-3
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { Application, AssistantResponse, createApp, Message } from '@callstack/byorg-core';
22
import type { CommandModule } from 'yargs';
33
import { checkIfConfigExists, parseConfigFile } from '../../config-file.js';
4-
import { getVerbose, output, outputError, setVerbose } from '../../output.js';
4+
import { getVerbose, output, outputError, outputSystem, setVerbose } from '../../output.js';
55
import { run as runInit } from '../init.js';
66
import { colorAssistant, colorVerbose } from '../../colors.js';
77
import { formatSpeed, formatTime } from '../../format.js';
@@ -10,7 +10,7 @@ import { processChatCommand } from './commands.js';
1010
import { cliOptions, type CliOptions } from './cli-options.js';
1111
import { getProvider, getProviderConfig, initProvider } from './providers.js';
1212
import { streamingClear, streamingFinish, streamingStart, streamingUpdate } from './streaming.js';
13-
import { messages } from './state.js';
13+
import { messages, updateUsage } from './state.js';
1414
import { texts } from './texts.js';
1515
import { exit } from './utils.js';
1616

@@ -54,6 +54,8 @@ async function run(initialPrompt: string, options: CliOptions) {
5454
await processMessages(app, messages);
5555
}
5656

57+
outputSystem(texts.initialHelp);
58+
5759
// eslint-disable-next-line no-constant-condition
5860
while (true) {
5961
const userMessage = await readUserInput();
@@ -83,8 +85,9 @@ async function processMessages(app: Application, messages: Message[]) {
8385
});
8486

8587
if (response.role === 'assistant') {
86-
messages.push({ role: 'assistant', content: response.content });
8788
streamingFinish(`${formatResponse(response)}\n`);
89+
messages.push({ role: 'assistant', content: response.content });
90+
updateUsage(response.usage);
8891
} else {
8992
streamingFinish(response.content);
9093
}

‎src/commands/chat/providers.tsx

+3-3
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import openAi from '../../engine/providers/open-ai.js';
1010
import anthropic from '../../engine/providers/anthropic.js';
1111
import perplexity from '../../engine/providers/perplexity.js';
1212
import mistral from '../../engine/providers/mistral.js';
13-
import { output, outputVerbose, outputWarning } from '../../output.js';
13+
import { outputSystem, outputVerbose, outputWarning } from '../../output.js';
1414
import { CliOptions } from './cli-options.js';
1515
import { filterOutApiKey, handleInputFile } from './utils.js';
1616

@@ -97,13 +97,13 @@ export function initProvider(options: CliOptions, configFile: ConfigFile) {
9797
systemPrompt: fileSystemPrompt,
9898
costWarning,
9999
costInfo,
100-
} = handleInputFile(options.file, getProviderConfig(), provider);
100+
} = handleInputFile(options.file, providerConfig, provider);
101101

102102
providerConfig.systemPrompt += `\n\n${fileSystemPrompt}`;
103103
if (costWarning) {
104104
outputWarning(costWarning);
105105
} else if (costInfo) {
106-
output(costInfo);
106+
outputSystem(costInfo);
107107
}
108108
}
109109
}

‎src/commands/chat/state.ts

+25-1
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,27 @@
1-
import { type Message } from '@callstack/byorg-core';
1+
import { ModelUsage, type Message } from '@callstack/byorg-core';
22

33
export const messages: Message[] = [];
4+
5+
export const totalUsage: ModelUsage = {
6+
inputTokens: 0,
7+
outputTokens: 0,
8+
requests: 0,
9+
responseTime: 0,
10+
model: '',
11+
usedTools: {},
12+
};
13+
14+
export function resetUsage() {
15+
totalUsage.inputTokens = 0;
16+
totalUsage.outputTokens = 0;
17+
totalUsage.requests = 0;
18+
totalUsage.responseTime = 0;
19+
}
20+
21+
export function updateUsage(usage: ModelUsage) {
22+
totalUsage.inputTokens += usage.inputTokens;
23+
totalUsage.outputTokens += usage.outputTokens;
24+
totalUsage.requests += usage.requests;
25+
totalUsage.responseTime += usage.responseTime;
26+
totalUsage.model = usage.model;
27+
}

‎src/commands/chat/texts.ts

+1-2
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
export const texts = {
22
userLabel: 'me:',
33
assistantLabel: 'ai:',
4-
initialHelp: 'Type "/exit" or press Ctrl+C to exit. Type "/help" to see available commands.',
5-
responseLoading: 'Thinking ...',
4+
initialHelp: 'Type "/exit" or press Ctrl+C to exit. Type "/help" to see available commands.\n',
65
} as const;

‎src/commands/chat/usage.ts

+27
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
import { ModelUsage } from '@callstack/byorg-core';
2+
import { Provider } from '../../engine/providers/provider.js';
3+
import { ProviderConfig } from '../../engine/providers/config.js';
4+
5+
export type CostContext = {
6+
provider: Provider;
7+
providerConfig: ProviderConfig;
8+
};
9+
10+
export function calculateUsageCost(usage: Partial<ModelUsage>, context: CostContext) {
11+
const pricing = getModelPricing(usage, context);
12+
if (pricing === undefined) {
13+
return undefined;
14+
}
15+
16+
const inputCost = ((usage.inputTokens ?? 0) * (pricing.inputTokensCost ?? 0)) / 1_000_000;
17+
const outputCost = ((usage.outputTokens ?? 0) * (pricing.outputTokensCost ?? 0)) / 1_000_000;
18+
const requestsCost = ((usage.requests ?? 0) * (pricing.requestsCost ?? 0)) / 1_000_000;
19+
return inputCost + outputCost + requestsCost;
20+
}
21+
22+
function getModelPricing(usage: Partial<ModelUsage>, { provider, providerConfig }: CostContext) {
23+
return (
24+
provider.modelPricing[usage.model ?? providerConfig.model] ??
25+
provider.modelPricing[providerConfig.model]
26+
);
27+
}

‎src/commands/chat/utils.ts

+3-7
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@ import {
77
FILE_COST_WARNING,
88
FILE_TOKEN_COUNT_WARNING,
99
} from '../../default-config.js';
10-
import { calculateUsageCost } from '../../engine/session.js';
1110
import { getTokensCount } from '../../engine/tokenizer.js';
1211
import type { ProviderConfig } from '../../engine/providers/config.js';
1312
import type { Provider } from '../../engine/providers/provider.js';
@@ -17,13 +16,12 @@ import {
1716
getDefaultFilename,
1817
getUniqueFilename,
1918
} from '../../file-utils.js';
20-
import { output } from '../../output.js';
2119
import { texts } from './texts.js';
2220
import { closeInput } from './input.js';
21+
import { calculateUsageCost } from './usage.js';
2322

2423
export function exit() {
2524
closeInput();
26-
output('\nBye...');
2725
process.exit(0);
2826
}
2927

@@ -35,7 +33,7 @@ interface HandleInputFileResult {
3533

3634
export function handleInputFile(
3735
inputFile: string,
38-
config: ProviderConfig,
36+
providerConfig: ProviderConfig,
3937
provider: Provider,
4038
): HandleInputFileResult {
4139
const filePath = path.resolve(inputFile.replace('~', os.homedir()));
@@ -46,9 +44,7 @@ export function handleInputFile(
4644

4745
const fileContent = fs.readFileSync(filePath).toString();
4846
const fileTokens = getTokensCount(fileContent);
49-
50-
const pricing = provider.modelPricing[config.model];
51-
const fileCost = calculateUsageCost({ inputTokens: fileTokens }, pricing);
47+
const fileCost = calculateUsageCost({ inputTokens: fileTokens }, { provider, providerConfig });
5248

5349
let costWarning = null;
5450
let costInfo = null;

‎src/engine/session.ts

-42
This file was deleted.

‎src/format.ts

+1-22
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,3 @@
1-
import type { SessionCost, SessionUsage } from './engine/session.js';
2-
31
export function formatCost(value: number | undefined, precision = 4) {
42
if (value == null) {
53
return '?';
@@ -29,32 +27,13 @@ export function formatTokenCount(tokenCount: number, roundTo = 1) {
2927
return `${scaledCount.toFixed(0)}${suffixes[suffixIndex]}`;
3028
}
3129

32-
export function formatSessionStats(responseTime?: number, usage?: SessionUsage) {
33-
const parts = [
34-
responseTime ? `time: ${(responseTime / 1000).toFixed(1)} s` : undefined,
35-
usage
36-
? `tokens: ${usage.current.inputTokens}+${usage.current.outputTokens} (total: ${usage.total.inputTokens}+${usage.total.outputTokens})`
37-
: undefined,
38-
];
39-
40-
return parts.filter((x) => x !== undefined).join(', ');
41-
}
42-
43-
export function formatSessionCost(cost: SessionCost | undefined) {
44-
if (cost === undefined) {
45-
return undefined;
46-
}
47-
48-
return `costs: ${formatCost(cost.current)} (total: ${formatCost(cost.total)})`;
49-
}
50-
5130
export function formatTime(timeInMs: number) {
5231
return `${(timeInMs / 1000).toFixed(1)} s`;
5332
}
5433

5534
export function formatSpeed(tokens: number, timeInMs: number) {
5635
if (tokens == null || timeInMs == null || timeInMs === 0) {
57-
return '? tok/s';
36+
return '? tokens/s';
5837
}
5938

6039
return `${((tokens * 1000) / timeInMs).toFixed(1)} tok/s`;

‎src/output.ts

+5-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { colorError, colorVerbose, colorWarning } from './colors.js';
1+
import { colorError, colorSystem, colorVerbose, colorWarning } from './colors.js';
22

33
let showVerbose = false;
44

@@ -14,6 +14,10 @@ export function output(text: string, ...args: unknown[]) {
1414
console.log(text, ...args);
1515
}
1616

17+
export function outputSystem(text: string, ...args: unknown[]) {
18+
console.log(colorSystem(text, ...args));
19+
}
20+
1721
export function outputVerbose(message: string, ...args: unknown[]) {
1822
if (showVerbose) {
1923
console.log(colorVerbose(message, ...args));

0 commit comments

Comments
 (0)