feat: record model parameters on generations

src/index.ts

@@ -264,6 +264,82 @@ export function toMultimodalContent(
   ];
 }
 
+const MAX_TOKENS_KEYS = ["max_tokens", "max_completion_tokens", "max_output_tokens", "maxOutputTokens", "maxTokens"];
+const MAX_PARAM_CHARS = 200;
+const THINKING_BUDGET_KEYS = [
+  "thinking_token_budget",
+  "thinking_budget",
+  "thinking_budget_tokens",
+  "budget_tokens",
+  "thinkingBudget",
+];
+
+function findNumber(bag: unknown, keys: string[], depth = 2): number | undefined {
+  if (!bag || typeof bag !== "object" || Array.isArray(bag)) return undefined;
+  const record = bag as Record<string, unknown>;
+  for (const key of keys) {
+    const value = record[key];
+    if (typeof value === "number" && Number.isInteger(value) && value > 0) return value;
+  }
+  if (depth === 0) return undefined;
+  for (const value of Object.values(record)) {
+    const found = findNumber(value, keys, depth - 1);
+    if (found !== undefined) return found;
+  }
+  return undefined;
+}
+
+function pickCacheRetention(payload: unknown): string | undefined {
+  const bag = payload as Record<string, unknown> | undefined;
+  const flat = bag?.prompt_cache_retention;
+  if (typeof flat === "string") return flat;
+  const system = bag?.system;
+  if (!Array.isArray(system)) return undefined;
+  for (const block of system) {
+    const ttl = (block as { cache_control?: { ttl?: unknown } } | undefined)?.cache_control?.ttl;
+    if (typeof ttl === "string") return ttl;
+  }
+  return undefined;
+}
+
+function pickToolChoice(payload: unknown): string | undefined {
+  const value = (payload as Record<string, unknown> | undefined)?.tool_choice;
+  if (typeof value === "string") return value;
+  const type = (value as { type?: unknown } | undefined)?.type;
+  return typeof type === "string" ? type : undefined;
+}
+
+export function extractModelParameters(
+  payload: unknown,
+  model: { reasoning?: boolean; samplingParams?: Record<string, unknown> } | undefined,
+  thinkingLevel?: string,
+): Record<string, string | number> | undefined {
+  const out: Record<string, string | number> = {};
+  try {
+    const maxTokens = findNumber(payload, MAX_TOKENS_KEYS);
+    if (maxTokens !== undefined) out.max_tokens = maxTokens;
+    if (model?.reasoning && thinkingLevel && thinkingLevel !== "off") out.thinking_level = thinkingLevel;
+    const thinkingBudget = findNumber(payload, THINKING_BUDGET_KEYS);
+    if (thinkingBudget !== undefined) out.thinking_budget_tokens = thinkingBudget;
+    const cacheRetention = pickCacheRetention(payload);
+    if (cacheRetention !== undefined) out.prompt_cache_retention = cacheRetention;
+    const serviceTier = (payload as Record<string, unknown> | undefined)?.service_tier;
+    if (typeof serviceTier === "string") out.service_tier = serviceTier;
+    const toolChoice = pickToolChoice(payload);
+    if (toolChoice !== undefined) out.tool_choice = toolChoice;
+    for (const key of Object.keys(model?.samplingParams ?? {})) {
+      const value = (payload as Record<string, unknown> | undefined)?.[key];
+      if (value === undefined) continue;
+      if (typeof value === "number") out[key] = value;
+      else {
+        const text = typeof value === "string" ? value : JSON.stringify(value);
+        if (text !== undefined && text.length <= MAX_PARAM_CHARS) out[key] = text;
+      }
+    }
+  } catch {}
+  return Object.keys(out).length ? out : undefined;
+}
+
 type PiAgentMessages = Parameters<typeof convertToLlm>[0];
 
 export interface ChatMlToolCall {
@@ -759,7 +835,7 @@ export default function (pi: ExtensionAPI) {
     debug("system prompt captured", systemPrompt.length);
   });
 
-  pi.on("before_provider_request", (_event, ctx) => {
+  pi.on("before_provider_request", (event, ctx) => {
     if (!state) return;
     if (state.openGeneration && !state.openGeneration.finished) {
       const gen = state.openGeneration;
@@ -787,11 +863,12 @@ export default function (pi: ExtensionAPI) {
       {
         input: generationInput,
         model: ctx.model?.id,
+        modelParameters: extractModelParameters(event.payload, ctx.model, ctx.thinkingLevel),
         metadata: {
           assistant_index: index - 1,
           input_source: history ? "context" : "delta",
           ...(history ? { history_message_count: history.length } : {}),
-          ...(ctx.model ? { provider: ctx.model.provider } : {}),
+          ...(ctx.model ? { provider: ctx.model.provider, context_window: ctx.model.contextWindow } : {}),
         },
       },
       { asType: "generation" },

test/helpers.ts

@@ -38,6 +38,8 @@ export interface MockProvider {
   port: number;
   close: () => void;
   sentMessages: OpenAiMessage[][];
+  /** Request bodies the provider received, in call order. */
+  payloads: () => Array<Record<string, unknown>>;
 }
 
 function writeSseEvent(res: http.ServerResponse, obj: unknown) {
@@ -101,6 +103,7 @@ export const FINAL_ANSWER_THINKING =
 
 export function startMockProvider(): Promise<MockProvider> {
   const sentMessages: OpenAiMessage[][] = [];
+  const received: Array<Record<string, unknown>> = [];
   const server = http.createServer((req, res) => {
     let body = "";
     req.on("data", (c) => (body += c));
@@ -110,6 +113,7 @@ export function startMockProvider(): Promise<MockProvider> {
         messages?: OpenAiMessage[];
         tools?: Array<{ function?: { name?: string } }>;
       };
+      received.push(payload as Record<string, unknown>);
       const messages = payload.messages ?? [];
       sentMessages.push(messages);
       const model = payload.model ?? "mock-gpt-1";
@@ -196,7 +200,7 @@ export function startMockProvider(): Promise<MockProvider> {
   return new Promise((resolvePromise) => {
     server.listen(0, "127.0.0.1", () => {
       const port = (server.address() as { port: number }).port;
-      resolvePromise({ port, sentMessages, close: () => server.close() });
+      resolvePromise({ port, sentMessages, close: () => server.close(), payloads: () => [...received] });
     });
   });
 }