Prompt reorg for caching + usage collection

docs/src/content/docs/reference/scripts/context.md

@@ -106,7 +106,7 @@
 ### Referencing
 
 The `def` function returns a variable name that can be used in the prompt.
-The name might be formatted diferently to accommodate the model's preference.
+The name might be formatted differently to accommodate the model's preference.
 
 ```js "const f = "
 const f = def("FILE", file)
@@ -182,6 +182,15 @@
 def("FILE", env.files, { sliceSample: 100 })
 ```
 
+### Prompt Caching
+
+You can specify `ephemeral: true` to turn on some prompt caching optimization. In paricular, a `def` with `ephemeral` will be rendered at the back of the prompt
+to persist the [cache prefix](https://openai.com/index/api-prompt-caching/).
+
+```js
+def("FILE", env.files, { ephemeral: true })
+```
+
 ## Data definition (`defData`)
 
 The `defData` function offers additional formatting options for converting a data object into a textual representation. It supports rendering objects as YAML, JSON, or CSV (formatted as a markdown table).

packages/core/src/chat.ts

@@ -28,6 +28,8 @@ import {
     ChatCompletionResponse,
     ChatCompletionsOptions,
     ChatCompletionTool,
+    ChatCompletionUsage,
+    ChatCompletionUsages,
     ChatCompletionUserMessageParam,
     CreateChatCompletionRequest,
 } from "./chattypes"
@@ -364,6 +366,7 @@ function structurifyChatSession(
     schemas: Record<string, JSONSchema>,
     genVars: Record<string, string>,
     options: GenerationOptions,
+    usages: ChatCompletionUsages,
     others?: {
         resp?: ChatCompletionResponse
         err?: any
@@ -426,17 +429,20 @@ function structurifyChatSession(
         error,
         genVars,
         schemas,
+        usages,
     }
 }
 
 async function processChatMessage(
+    req: CreateChatCompletionRequest,
     resp: ChatCompletionResponse,
     messages: ChatCompletionMessageParam[],
     tools: ToolCallback[],
     chatParticipants: ChatParticipant[],
     schemas: Record<string, JSONSchema>,
     genVars: Record<string, string>,
-    options: GenerationOptions
+    options: GenerationOptions,
+    usages: ChatCompletionUsages
 ): Promise<RunPromptResult> {
     const {
         stats,
@@ -445,6 +451,8 @@ async function processChatMessage(
         cancellationToken,
     } = options
 
+    accumulateChatUsage(usages, req.model, resp.usage)
+
     if (resp.text)
         messages.push({
             role: "assistant",
@@ -534,11 +542,29 @@ export function mergeGenerationOptions(
     }
 }
 
+function accumulateChatUsage(
+    usages: ChatCompletionUsages,
+    model: string,
+    usage: ChatCompletionUsage
+) {
+    if (!usage) return
+
+    const u =
+        usages[model] ??
+        (usages[model] = <ChatCompletionUsage>{
+            completion_tokens: 0,
+            prompt_tokens: 0,
+            total_tokens: 0,
+        })
+    u.completion_tokens += u.completion_tokens
+    u.prompt_tokens += u.prompt_tokens
+    u.total_tokens += u.total_tokens
+}
+
 export async function executeChatSession(
     connectionToken: LanguageModelConfiguration,
     cancellationToken: CancellationToken,
     messages: ChatCompletionMessageParam[],
-    vars: Partial<ExpansionVariables>,
     toolDefinitions: ToolCallback[],
     schemas: Record<string, JSONSchema>,
     completer: ChatCompletionHandler,
@@ -567,6 +593,7 @@ export async function executeChatSession(
         : undefined
     trace.startDetails(`🧠 llm chat`)
     if (tools?.length) trace.detailsFenced(`🛠️ tools`, tools, "yaml")
+    const usages: ChatCompletionUsages = {}
     try {
         let genVars: Record<string, string>
         while (true) {
@@ -585,34 +612,35 @@ export async function executeChatSession(
             let resp: ChatCompletionResponse
             try {
                 checkCancelled(cancellationToken)
+                const req: CreateChatCompletionRequest = {
+                    model,
+                    temperature: temperature,
+                    top_p: topP,
+                    max_tokens: maxTokens,
+                    seed,
+                    stream: true,
+                    messages,
+                    tools,
+                    response_format:
+                        responseType === "json_object"
+                            ? { type: responseType }
+                            : responseType === "json_schema"
+                              ? {
+                                    type: "json_schema",
+                                    json_schema: {
+                                        name: "result",
+                                        schema: toStrictJSONSchema(
+                                            responseSchema
+                                        ),
+                                        strict: true,
+                                    },
+                                }
+                              : undefined,
+                }
                 try {
                     trace.startDetails(`📤 llm request`)
                     resp = await completer(
-                        {
-                            model,
-                            temperature: temperature,
-                            top_p: topP,
-                            max_tokens: maxTokens,
-                            seed,
-                            stream: true,
-                            messages,
-                            tools,
-                            response_format:
-                                responseType === "json_object"
-                                    ? { type: responseType }
-                                    : responseType === "json_schema"
-                                      ? {
-                                            type: "json_schema",
-                                            json_schema: {
-                                                name: "result",
-                                                schema: toStrictJSONSchema(
-                                                    responseSchema
-                                                ),
-                                                strict: true,
-                                            },
-                                        }
-                                      : undefined,
-                        },
+                        req,
                         connectionToken,
                         genOptions,
                         trace
@@ -625,13 +653,15 @@ export async function executeChatSession(
                 }
 
                 const output = await processChatMessage(
+                    req,
                     resp,
                     messages,
                     toolDefinitions,
                     chatParticipants,
                     schemas,
                     genVars,
-                    genOptions
+                    genOptions,
+                    usages
                 )
                 if (output) return output
             } catch (err) {
@@ -640,6 +670,7 @@ export async function executeChatSession(
                     schemas,
                     genVars,
                     genOptions,
+                    usages,
                     { resp, err }
                 )
             }

packages/core/src/chattypes.ts

@@ -18,6 +18,15 @@ export interface AICIRequest {
 }
 
 // Aliases for OpenAI chat completion types
+export type ChatCompletionUsage = Omit<
+    OpenAI.Completions.CompletionUsage,
+    "completion_tokens_details"
+>
+
+/**
+ * Per model storage of chat completion usages.
+ */
+export type ChatCompletionUsages = Record<string, ChatCompletionUsage>
 
 // Text content part of a chat completion
 export type ChatCompletionContentPartText =
@@ -99,6 +108,7 @@ export interface ChatCompletionResponse {
     toolCalls?: ChatCompletionToolCall[] // List of tool calls made during the response
     finishReason?: // Reason why the chat completion finished
     "stop" | "length" | "tool_calls" | "content_filter" | "cancel" | "fail"
+    usage?: ChatCompletionUsage // Usage information for the completion
 }
 
 // Alias for OpenAI's API error type

packages/core/src/openai.ts

@@ -1,4 +1,4 @@
-import { normalizeInt, trimTrailingSlash } from "./util"
+import { logVerbose, normalizeInt, trimTrailingSlash } from "./util"
 import { LanguageModelConfiguration, host } from "./host"
 import {
     AZURE_OPENAI_API_VERSION,
@@ -19,6 +19,7 @@ import {
     ChatCompletionToolCall,
     ChatCompletionResponse,
     ChatCompletionChunk,
+    ChatCompletionUsage,
 } from "./chattypes"
 import { resolveTokenEncoder } from "./encoders"
 import { toSignal } from "./cancellation"
@@ -93,17 +94,20 @@ export const OpenAIChatCompletion: ChatCompletionHandler = async (
         return { text: cached, finishReason: cachedFinishReason, cached: true }
     }
 
-    const r2 = { ...req, model }
+    const r2 = {
+        ...req,
+        stream: true,
+        stream_options: { include_usage: true },
+        model,
+    }
     let postReq: any = r2
 
     let url = ""
     const toolCalls: ChatCompletionToolCall[] = []
 
     if (cfg.type === "openai" || cfg.type === "localai") {
-        r2.stream = true
         url = trimTrailingSlash(cfg.base) + "/chat/completions"
     } else if (cfg.type === "azure") {
-        r2.stream = true
         delete r2.model
         url =
             trimTrailingSlash(cfg.base) +
@@ -175,6 +179,7 @@ export const OpenAIChatCompletion: ChatCompletionHandler = async (
     let finishReason: ChatCompletionResponse["finishReason"] = undefined
     let chatResp = ""
     let pref = ""
+    let usage: ChatCompletionUsage
 
     const decoder = host.createUTF8Decoder()
     if (r.body.getReader) {
@@ -193,15 +198,22 @@ export const OpenAIChatCompletion: ChatCompletionHandler = async (
     if (cancellationToken?.isCancellationRequested) finishReason = "cancel"
 
     trace.appendContent("\n\n")
-    trace.itemValue(`finish reason`, finishReason)
+    trace.itemValue(`🏁 finish reason`, finishReason)
+    if (usage) {
+        trace.itemValue(
+            `🪙 tokens`,
+            `${usage.total_tokens} total, ${usage.prompt_tokens} prompt, ${usage.completion_tokens} completion`
+        )
+    }
+
     if (done && finishReason === "stop")
         await cacheStore.set(
             cachedKey,
             { text: chatResp, finishReason },
             { trace }
         )
 
-    return { text: chatResp, toolCalls, finishReason }
+    return { text: chatResp, toolCalls, finishReason, usage }
 
     function doChunk(value: Uint8Array) {
         // Massage and parse the chunk of data
@@ -216,6 +228,7 @@ export const OpenAIChatCompletion: ChatCompletionHandler = async (
             }
             try {
                 const obj: ChatCompletionChunk = JSON.parse(json)
+                if (obj.usage) usage = obj.usage
                 if (!obj.choices?.length) return ""
                 else if (obj.choices?.length != 1)
                     throw new Error("too many choices in response")

packages/core/src/promptcontext.ts

@@ -234,7 +234,7 @@ export async function createPromptContext(
     })
 
     // Freeze project options to prevent modification
-    const projectOptions = Object.freeze({ prj, vars, env })
+    const projectOptions = Object.freeze({ prj, env })
     const ctx: PromptContext & RunPromptContextNode = {
         ...createChatGenerationContext(options, trace, projectOptions),
         script: () => {},