diff --git a/packages/cli/src/run.ts b/packages/cli/src/run.ts index 4c446c4b8f..9805fa3a02 100644 --- a/packages/cli/src/run.ts +++ b/packages/cli/src/run.ts @@ -1,7 +1,7 @@ import { capitalize } from "inflection" import { resolve, join, relative, dirname } from "node:path" import { isQuiet, wrapColor } from "./log" -import { emptyDir, ensureDir, appendFileSync } from "fs-extra" +import { emptyDir, ensureDir, appendFileSync, exists } from "fs-extra" import { convertDiagnosticsToSARIF } from "./sarif" import { buildProject } from "./build" import { diagnosticsToCSV } from "../../core/src/ast" @@ -34,6 +34,7 @@ import { TRACE_DETAILS, CLI_ENV_VAR_RX, AGENT_MEMORY_CACHE_NAME, + STATS_DIR_NAME, } from "../../core/src/constants" import { isCancelError, errorMessage } from "../../core/src/error" import { Fragment, GenerationResult } from "../../core/src/generation" @@ -74,6 +75,8 @@ import { prettifyMarkdown } from "../../core/src/markdown" import { delay } from "es-toolkit" import { GenerationStats } from "../../core/src/usage" import { traceAgentMemory } from "../../core/src/agent" +import { JSONLineCache } from "../../core/src/cache" +import { appendFile, stat } from "node:fs/promises" function parseVars( vars: string[], @@ -340,6 +343,7 @@ export async function runScript( } if (!isQuiet) logVerbose("") // force new line + await aggregateResults(scriptId, outTrace, result) await traceAgentMemory(trace) if (outAnnotations && result.annotations?.length) { if (isJSONLFilename(outAnnotations)) @@ -544,3 +548,44 @@ export async function runScript( return { exitCode: 0, result } } + +async function aggregateResults( + scriptId: string, + outTrace: string, + result: GenerationResult +) { + const statsDir = dotGenaiscriptPath(".") + await ensureDir(statsDir) + const statsFile = path.join(statsDir, "stats.csv") + if (!(await exists(statsFile))) + await writeFile( + statsFile, + [ + "script", + "status", + "cost", + "total_tokens", + "prompt_tokens", + "completion_tokens", + "trace", + "version", + ].join(",") + "\n", + { encoding: "utf-8" } + ) + await appendFile( + statsFile, + [ + scriptId, + result.status, + result.stats.cost, + result.stats.total_tokens, + result.stats.prompt_tokens, + result.stats.completion_tokens, + path.basename(outTrace), + result.version, + ] + .map((s) => String(s)) + .join(",") + "\n", + { encoding: "utf-8" } + ) +} diff --git a/packages/core/src/constants.ts b/packages/core/src/constants.ts index 89af6d20a9..bec63aee27 100644 --- a/packages/core/src/constants.ts +++ b/packages/core/src/constants.ts @@ -120,6 +120,7 @@ export const PROMPTFOO_CONFIG_DIR = ".genaiscript/config/tests" export const PROMPTFOO_REMOTE_API_PORT = 15500 export const RUNS_DIR_NAME = "runs" +export const STATS_DIR_NAME = "stats" export const EMOJI_SUCCESS = "✅" export const EMOJI_FAIL = "❌" diff --git a/packages/core/src/generation.ts b/packages/core/src/generation.ts index 532a661851..041eac719d 100644 --- a/packages/core/src/generation.ts +++ b/packages/core/src/generation.ts @@ -1,7 +1,11 @@ // Import necessary modules and interfaces import { CancellationToken } from "./cancellation" import { LanguageModel } from "./chat" -import { ChatCompletionMessageParam, ChatCompletionsOptions } from "./chattypes" +import { + ChatCompletionMessageParam, + ChatCompletionsOptions, + ChatCompletionUsage, +} from "./chattypes" import { MarkdownTrace } from "./trace" import { GenerationStats } from "./usage" @@ -66,6 +70,13 @@ export interface GenerationResult extends GenerationOutput { * Version of the GenAIScript used */ version: string + + /** + * Statistics of the generation + */ + stats: { + cost: number + } & ChatCompletionUsage } // Type representing possible statuses of generation diff --git a/packages/core/src/promptrunner.ts b/packages/core/src/promptrunner.ts index 6b348cf19c..06b3f0d9de 100644 --- a/packages/core/src/promptrunner.ts +++ b/packages/core/src/promptrunner.ts @@ -275,6 +275,10 @@ export async function runTemplate( genVars, schemas, json, + stats: { + cost: options.stats.cost(), + ...options.stats.usage, + }, } // If there's an error, provide status text diff --git a/packages/core/src/usage.ts b/packages/core/src/usage.ts index 5761ac3a35..8896c9dc9a 100644 --- a/packages/core/src/usage.ts +++ b/packages/core/src/usage.ts @@ -20,7 +20,7 @@ import { /** * Estimates the cost of a chat completion based on the model and usage. - * + * * @param modelId - The identifier of the model used for chat completion. * @param usage - The usage statistics for the chat completion. * @returns The estimated cost or undefined if estimation is not possible. @@ -57,7 +57,7 @@ export function estimateCost(modelId: string, usage: ChatCompletionUsage) { /** * Renders the cost as a string for display purposes. - * + * * @param value - The cost to be rendered. * @returns A string representation of the cost. */ @@ -89,7 +89,7 @@ export class GenerationStats { /** * Constructs a GenerationStats instance. - * + * * @param model - The model used for chat completions. * @param label - Optional label for the statistics. */ @@ -113,7 +113,7 @@ export class GenerationStats { /** * Calculates the total cost based on the usage statistics. - * + * * @returns The total cost. */ cost(): number { @@ -125,7 +125,7 @@ export class GenerationStats { /** * Accumulates the usage statistics from this instance and its children. - * + * * @returns The accumulated usage statistics. */ accumulatedUsage(): ChatCompletionUsage { @@ -149,7 +149,7 @@ export class GenerationStats { /** * Creates a new child GenerationStats instance. - * + * * @param model - The model used for the child chat completions. * @param label - Optional label for the child's statistics. * @returns The created child GenerationStats instance. @@ -162,7 +162,7 @@ export class GenerationStats { /** * Traces the generation statistics using a MarkdownTrace instance. - * + * * @param trace - The MarkdownTrace instance used for tracing. */ trace(trace: MarkdownTrace) { @@ -176,7 +176,7 @@ export class GenerationStats { /** * Helper method to trace individual statistics. - * + * * @param trace - The MarkdownTrace instance used for tracing. */ private traceStats(trace: MarkdownTrace) { @@ -218,7 +218,7 @@ export class GenerationStats { /** * Helper method to log tokens with indentation. - * + * * @param indent - The indentation used for logging. */ private logTokens(indent: string) { @@ -240,7 +240,7 @@ export class GenerationStats { /** * Adds usage statistics to the current instance. - * + * * @param req - The request containing details about the chat completion. * @param usage - The usage statistics to be added. */