Skip to content

Commit

Permalink
feat(logger): add token usage details in summary (#255)
Browse files Browse the repository at this point in the history
  • Loading branch information
@khalatevarun
khalatevarun authored Jan 4, 2025
1 parent bec0434 commit 61ca4a7
Show file tree
Hide file tree
Showing 3 changed files with 102 additions and 8 deletions.
26 changes: 23 additions & 3 deletions packages/shortest/src/ai/client.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,11 @@ export class AIClient {
content: Anthropic.Beta.Messages.BetaContentBlockParam,
) => void,
toolOutputCallback?: (name: string, input: any) => void,
) {
): Promise<{
finalResponse: any;
tokenUsage: { input: number; output: number };
pendingCache: any;
}> {
const maxRetries = 3;
let attempts = 0;

Expand All @@ -56,6 +60,11 @@ export class AIClient {
await new Promise((r) => setTimeout(r, 5000 * attempts));
}
}
return {
finalResponse: null,
tokenUsage: { input: 0, output: 0 },
pendingCache: null,
};
}

async makeRequest(
Expand All @@ -65,7 +74,12 @@ export class AIClient {
content: Anthropic.Beta.Messages.BetaContentBlockParam,
) => void,
_toolOutputCallback?: (name: string, input: any) => void,
) {
): Promise<{
messages: any;
finalResponse: any;
pendingCache: any;
tokenUsage: { input: number; output: number };
}> {
const messages: Anthropic.Beta.Messages.BetaMessageParam[] = [];
// temp cache store
const pendingCache: Partial<{ steps?: CacheStep[] }> = {};
Expand All @@ -92,8 +106,13 @@ export class AIClient {
tools: [...AITools],
betas: ["computer-use-2024-10-22"],
});

// Log AI response and tool usage

const tokenUsage = {
input: response.usage.input_tokens,
output: response.usage.output_tokens,
};

if (this.debugMode) {
response.content.forEach((block) => {
if (block.type === "text") {
Expand Down Expand Up @@ -232,6 +251,7 @@ export class AIClient {
messages,
finalResponse: response,
pendingCache,
tokenUsage,
};
}
} catch (error: any) {
Expand Down
19 changes: 15 additions & 4 deletions packages/shortest/src/core/runner/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -140,7 +140,11 @@ export class TestRunner {
test: TestFunction,
context: BrowserContext,
config: { noCache: boolean } = { noCache: false },
) {
): Promise<{
result: "pass" | "fail";
reason: string;
tokenUsage: { input: number; output: number };
}> {
// If it's direct execution, skip AI
if (test.directExecution) {
try {
Expand All @@ -149,12 +153,14 @@ export class TestRunner {
return {
result: "pass" as const,
reason: "Direct execution successful",
tokenUsage: { input: 0, output: 0 },
};
} catch (error) {
return {
result: "fail" as const,
reason:
error instanceof Error ? error.message : "Direct execution failed",
tokenUsage: { input: 0, output: 0 },
};
}
}
Expand Down Expand Up @@ -238,10 +244,11 @@ export class TestRunner {
: error instanceof Error
? error.message
: String(error),
tokenUsage: { input: 0, output: 0 },
};
}
}
return result;
return { ...result, tokenUsage: { input: 0, output: 0 } };
} catch {
// delete stale cached test entry
await this.cache.delete(test);
Expand All @@ -263,6 +270,7 @@ export class TestRunner {
return {
result: "fail" as const,
reason: error instanceof Error ? error.message : String(error),
tokenUsage: { input: 0, output: 0 },
};
}
}
Expand All @@ -276,7 +284,7 @@ export class TestRunner {

// Parse AI result first
const finalMessage = result.finalResponse.content.find(
(block) =>
(block: any) =>
block.type === "text" &&
(block as Anthropic.Beta.Messages.BetaTextBlock).text.includes(
'"result":',
Expand Down Expand Up @@ -311,6 +319,7 @@ export class TestRunner {
: error instanceof Error
? error.message
: String(error),
tokenUsage: result.tokenUsage,
};
}
}
Expand All @@ -319,7 +328,7 @@ export class TestRunner {
// batch set new chache if test is successful
await this.cache.set(test, result.pendingCache);
}
return aiResult;
return { ...aiResult, tokenUsage: result.tokenUsage };
}

private async executeTestFile(file: string) {
Expand Down Expand Up @@ -354,6 +363,7 @@ export class TestRunner {
test.name,
result.result === "pass" ? "passed" : "failed",
result.result === "fail" ? new Error(result.reason) : undefined,
result.tokenUsage,
);

// Execute afterEach hooks with shared context
Expand Down Expand Up @@ -435,6 +445,7 @@ export class TestRunner {
) {
// @ts-expect-error
const [x, y] = step.action.input.coordinate;

const componentStr =
await browserTool.getNormalizedComponentStringByCoords(x, y);

Expand Down
65 changes: 64 additions & 1 deletion packages/shortest/src/utils/logger.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,17 +3,27 @@ import { AssertionError } from "../types";

export type TestStatus = "pending" | "running" | "passed" | "failed";

interface TokenMetrics {
input: number;
output: number;
}

interface TestResult {
name: string;
status: TestStatus;
error?: Error;
tokenUsage?: TokenMetrics;
}

export class Logger {
private currentFile: string = "";
private testResults: TestResult[] = [];
private startTime: number = Date.now();

// token pricing (Claude 3.5 Sonnet)
private readonly COST_PER_1K_INPUT_TOKENS = 0.003;
private readonly COST_PER_1K_OUTPUT_TOKENS = 0.015;

startFile(file: string) {
this.currentFile = file.split("/").pop() || file;
console.log(pc.blue(`\n📄 ${pc.bold(this.currentFile)}`));
Expand All @@ -23,18 +33,60 @@ export class Logger {
name: string | undefined,
status: "passed" | "failed",
error?: Error,
tokenUsage?: TokenMetrics,
) {
const testName = name || "Unnamed Test";
const symbol = status === "passed" ? "✓" : "✗";
const color = status === "passed" ? pc.green : pc.red;

console.log(` ${color(symbol)} ${testName}`);

if (tokenUsage) {
const totalTokens = tokenUsage.input + tokenUsage.output;
const cost = this.calculateCost(tokenUsage.input, tokenUsage.output);
console.log(
pc.dim(
` ↳ ${totalTokens.toLocaleString()} tokens ` +
`(≈ $${cost.toFixed(2)})`,
),
);
}

if (error) {
console.log(pc.red(` ${error.message}`));
}

this.testResults.push({ name: testName, status, error });
this.testResults.push({ name: testName, status, error, tokenUsage });
}

private calculateCost(inputTokens: number, outputTokens: number): number {
const inputCost = (inputTokens / 1000) * this.COST_PER_1K_INPUT_TOKENS;
const outputCost = (outputTokens / 1000) * this.COST_PER_1K_OUTPUT_TOKENS;
return Number((inputCost + outputCost).toFixed(3));
}

private calculateTotalTokenUsage(): {
totalInputTokens: number;
totalOutputTokens: number;
totalCost: number;
} {
let totalInputTokens = 0;
let totalOutputTokens = 0;

this.testResults.forEach((result) => {
if (result.tokenUsage) {
totalInputTokens += result.tokenUsage.input;
totalOutputTokens += result.tokenUsage.output;
}
});

const totalCost = this.calculateCost(totalInputTokens, totalOutputTokens);

return {
totalInputTokens,
totalOutputTokens,
totalCost,
};
}

private getStatusIcon(status: TestStatus): string {
Expand All @@ -58,6 +110,10 @@ export class Logger {
).length;
const passedTests = totalTests - failedTests;

const { totalInputTokens, totalOutputTokens, totalCost } =
this.calculateTotalTokenUsage();
const totalTokens = totalInputTokens + totalOutputTokens;

console.log(pc.dim("⎯".repeat(50)));

console.log(
Expand All @@ -73,6 +129,13 @@ export class Logger {
pc.bold(" Start at "),
pc.dim(new Date(this.startTime).toLocaleTimeString()),
);
console.log(
pc.bold(" Tokens "),
pc.dim(
`${totalTokens.toLocaleString()} tokens ` +
`(≈ $${totalCost.toFixed(2)})`,
),
);
console.log(pc.dim("\n" + "⎯".repeat(50)));
}

Expand Down

0 comments on commit 61ca4a7

Please sign in to comment.