feat: 🎉 add stats aggregation for generation results

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" }
+    )
+}

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 = "❌"

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

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

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.
      */