Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: restore intro message and costs calculation #69

Merged
merged 2 commits into from
Dec 8, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions src/colors.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
62 changes: 37 additions & 25 deletions src/commands/chat/commands.ts
Original file line number Diff line number Diff line change
@@ -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('/')) {
Expand All @@ -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;
}

Expand All @@ -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);
}
9 changes: 6 additions & 3 deletions src/commands/chat/index.tsx
Original file line number Diff line number Diff line change
@@ -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';
Expand All @@ -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';

Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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);
}
Expand Down
6 changes: 3 additions & 3 deletions src/commands/chat/providers.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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';

Expand Down Expand Up @@ -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);
}
}
}
Expand Down
26 changes: 25 additions & 1 deletion src/commands/chat/state.ts
Original file line number Diff line number Diff line change
@@ -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;
}
3 changes: 1 addition & 2 deletions src/commands/chat/texts.ts
Original file line number Diff line number Diff line change
@@ -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;
27 changes: 27 additions & 0 deletions src/commands/chat/usage.ts
Original file line number Diff line number Diff line change
@@ -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]
);
}
10 changes: 3 additions & 7 deletions src/commands/chat/utils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -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);
}

Expand All @@ -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()));
Expand All @@ -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;
Expand Down
42 changes: 0 additions & 42 deletions src/engine/session.ts

This file was deleted.

23 changes: 1 addition & 22 deletions src/format.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,3 @@
import type { SessionCost, SessionUsage } from './engine/session.js';

export function formatCost(value: number | undefined, precision = 4) {
if (value == null) {
return '?';
Expand Down Expand Up @@ -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`;
Expand Down
6 changes: 5 additions & 1 deletion src/output.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import { colorError, colorVerbose, colorWarning } from './colors.js';
import { colorError, colorSystem, colorVerbose, colorWarning } from './colors.js';

let showVerbose = false;

Expand All @@ -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));
Expand Down
Loading