diff --git a/src/index.ts b/src/index.ts index 79f4e16..18868fa 100644 --- a/src/index.ts +++ b/src/index.ts @@ -291,6 +291,14 @@ 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; @@ -307,9 +315,29 @@ function findNumber(bag: unknown, keys: string[], depth = 2): number | undefined return undefined; } +function pickCacheRetention(payload: unknown): string | undefined { + const bag = payload as Record | 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 | 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 } | undefined, + model: { reasoning?: boolean; samplingParams?: Record } | undefined, thinkingLevel?: string, ): Record | undefined { const out: Record = {}; @@ -317,6 +345,23 @@ export function extractModelParameters( 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 | 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 | 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; } @@ -710,7 +755,7 @@ export default function (pi: ExtensionAPI) { modelParameters: extractModelParameters(event.payload, ctx.model, ctx.thinkingLevel), metadata: { assistant_index: index - 1, - ...(ctx.model ? { provider: ctx.model.provider } : {}), + ...(ctx.model ? { provider: ctx.model.provider, context_window: ctx.model.contextWindow } : {}), }, }, { asType: "generation" }, diff --git a/test/helpers.ts b/test/helpers.ts index 0be3873..e292310 100644 --- a/test/helpers.ts +++ b/test/helpers.ts @@ -75,7 +75,13 @@ function streamChunks( /** Usage the mock reports for a compaction/branch summarization call. */ export const SUMMARIZATION_USAGE = { prompt: 3571, completion: 313 }; -export function startMockProvider(): Promise<{ port: number; close: () => void }> { +export function startMockProvider(): Promise<{ + port: number; + close: () => void; + /** Request bodies the provider received, in call order. */ + payloads: () => Array>; +}> { + const received: Array> = []; const server = http.createServer((req, res) => { let body = ""; req.on("data", (c) => (body += c)); @@ -85,6 +91,7 @@ export function startMockProvider(): Promise<{ port: number; close: () => void } messages?: OpenAiMessage[]; tools?: Array<{ function?: { name?: string } }>; }; + received.push(payload as Record); const messages = payload.messages ?? []; const model = payload.model ?? "mock-gpt-1"; const lastUserIdx = messages.map((m) => m.role).lastIndexOf("user"); @@ -169,7 +176,7 @@ export function startMockProvider(): Promise<{ port: number; close: () => void } return new Promise((resolvePromise) => { server.listen(0, "127.0.0.1", () => { const port = (server.address() as { port: number }).port; - resolvePromise({ port, close: () => server.close() }); + resolvePromise({ port, close: () => server.close(), payloads: () => [...received] }); }); }); } diff --git a/test/model-parameters.test.ts b/test/model-parameters.test.ts index 58c4af6..bb4c182 100644 --- a/test/model-parameters.test.ts +++ b/test/model-parameters.test.ts @@ -16,6 +16,7 @@ import { describe("extractModelParameters", () => { const plain = { reasoning: false }; + const PARAM_CHAR_CAP = 200; it("reads the cap the payload really carries, whatever the dialect calls it", () => { const payloads: Array<[string, unknown]> = [ @@ -65,6 +66,107 @@ describe("extractModelParameters", () => { assert.deepEqual(extractModelParameters({}, { reasoning: true }, "medium"), { thinking_level: "medium" }); }); + it("reads the thinking budget the provider was really given", () => { + const reasoning = { reasoning: true }; + const payloads: Array<[string, unknown]> = [ + ["anthropic-messages", { max_tokens: 8192, thinking: { type: "enabled", budget_tokens: 7168 } }], + [ + "bedrock-converse-stream", + { + inferenceConfig: { maxTokens: 8192 }, + additionalModelRequestFields: { thinking: { type: "enabled", budget_tokens: 7168 } }, + }, + ], + ["google", { config: { maxOutputTokens: 8192, thinkingConfig: { thinkingBudget: 7168 } } }], + ["openai-completions (vLLM)", { max_tokens: 8192, thinking_token_budget: 7168 }], + ["openai-completions (thinking_budget)", { max_tokens: 8192, thinking_budget: 7168 }], + ["openai-completions (thinking_budget_tokens)", { max_tokens: 8192, thinking_budget_tokens: 7168 }], + ]; + for (const [dialect, payload] of payloads) { + assert.deepEqual( + extractModelParameters(payload, reasoning, "high"), + { max_tokens: 8192, thinking_level: "high", thinking_budget_tokens: 7168 }, + dialect, + ); + } + }); + + it("omits the budget when the provider gets none", () => { + assert.deepEqual( + extractModelParameters({ max_tokens: 8192, thinking: { type: "adaptive" }, output_config: { effort: "high" } }, { reasoning: true }, "high"), + { max_tokens: 8192, thinking_level: "high" }, + ); + assert.deepEqual(extractModelParameters({ config: { thinkingConfig: { thinkingBudget: 0 } } }, plain), undefined); + }); + + it("reports the prompt cache retention both dialects express differently", () => { + assert.deepEqual(extractModelParameters({ prompt_cache_retention: "24h" }, plain), { + prompt_cache_retention: "24h", + }); + assert.deepEqual( + extractModelParameters({ max_tokens: 8192, system: [{ type: "text", cache_control: { type: "ephemeral", ttl: "1h" } }] }, plain), + { max_tokens: 8192, prompt_cache_retention: "1h" }, + ); + assert.deepEqual( + extractModelParameters({ max_tokens: 8192, system: [{ type: "text", cache_control: { type: "ephemeral" } }] }, plain), + { max_tokens: 8192 }, + ); + }); + + it("reports the service tier and the tool choice when they are set", () => { + assert.deepEqual(extractModelParameters({ service_tier: "flex", tool_choice: "auto" }, plain), { + service_tier: "flex", + tool_choice: "auto", + }); + assert.deepEqual(extractModelParameters({ tool_choice: { type: "any" } }, plain), { tool_choice: "any" }); + assert.deepEqual(extractModelParameters({ tool_choice: { type: "function", function: { name: "read" } } }, plain), { + tool_choice: "function", + }); + }); + + it("reports the sampling parameters the model declares, with the values the wire carried", () => { + const samplingParams = { temperature: 0.2, top_p: 0.9, top_k: 40, min_p: 0.05, repetition_penalty: 1.1 }; + const onTheWire = { temperature: 0.7, top_p: 0.5, top_k: 10, min_p: 0.01, repetition_penalty: 1.9 }; + assert.deepEqual(extractModelParameters({ max_tokens: 8192, ...onTheWire }, { reasoning: false, samplingParams }), { + max_tokens: 8192, + ...onTheWire, + }); + }); + + it("serializes a non-scalar sampling parameter instead of dropping it", () => { + const samplingParams = { stop: [""], seed: 7 }; + assert.deepEqual(extractModelParameters({ ...samplingParams }, { reasoning: false, samplingParams }), { + stop: '[""]', + seed: 7, + }); + }); + + it("never lets a sampling parameter pull a payload structure onto the span", () => { + const conversation = [{ role: "user", content: "private user text ".repeat(40) }]; + for (const key of ["messages", "system", "tools", "input", "contents"]) { + const out = extractModelParameters( + { max_tokens: 8192, [key]: conversation }, + { reasoning: false, samplingParams: { [key]: 1 } }, + ); + assert.deepEqual(out, { max_tokens: 8192 }, key); + } + const long = "x".repeat(PARAM_CHAR_CAP + 1); + assert.deepEqual(extractModelParameters({ note: long }, { reasoning: false, samplingParams: { note: 1 } }), undefined); + const short = "x".repeat(PARAM_CHAR_CAP); + assert.deepEqual(extractModelParameters({ note: short }, { reasoning: false, samplingParams: { note: 1 } }), { + note: short, + }); + }); + + it("keeps the chips it already collected when a late sampling parameter blows up", () => { + const cyclic: Record = {}; + cyclic.self = cyclic; + assert.deepEqual( + extractModelParameters({ max_tokens: 1477, weird: cyclic }, { reasoning: true, samplingParams: { weird: 1 } }, "high"), + { max_tokens: 1477, thinking_level: "high" }, + ); + }); + it("rejects a cap that is not a positive whole number", () => { for (const bad of [Number.NaN, Number.POSITIVE_INFINITY, -5, 1.5, "8192", null, true]) { assert.equal(extractModelParameters({ max_tokens: bad }, plain), undefined, String(bad)); @@ -80,7 +182,13 @@ describe("extractModelParameters", () => { }, }); assert.equal(extractModelParameters(hostile, plain), undefined); - assert.equal(extractModelParameters(hostile, { reasoning: true }, "high"), undefined); + }); + + it("omits a declared sampling parameter the dialect never sent", () => { + assert.deepEqual( + extractModelParameters({ max_tokens: 8192 }, { reasoning: false, samplingParams: { temperature: 0.2 } }), + { max_tokens: 8192 }, + ); }); it("never descends into message arrays looking for a cap", () => { @@ -182,4 +290,81 @@ describe("integration: model parameters", () => { capture.close(); } }); + + it("records the long prompt cache retention pi put on the wire", async () => { + const capture = await startCaptureServer(); + try { + const sandbox = createSandbox(mock.port); + const modelsPath = join(sandbox.agentDir, "models.json"); + const models = JSON.parse(readFileSync(modelsPath, "utf8")); + models.providers.mock.models[0].compat = { supportsLongCacheRetention: true }; + writeFileSync(modelsPath, JSON.stringify(models)); + const result = await runPi(sandbox, "Explore this project and summarize it", { + env: { ...langfuseEnv(capture), PI_CACHE_RETENTION: "long" }, + }); + assert.equal(result.status, 0, `pi failed: ${result.stderr}`); + await waitForRequests(capture, 1); + + const params = generationParameters(capture.spans()); + assert.equal(params.length, 3); + for (const p of params) { + assert.deepEqual(p, { max_tokens: 8192, prompt_cache_retention: "24h" }); + } + } finally { + capture.close(); + } + }); + + it("reports exactly the sampling parameters the request carried", async () => { + const capture = await startCaptureServer(); + const wired = await startMockProvider(); + try { + const sandbox = createSandbox(wired.port); + const modelsPath = join(sandbox.agentDir, "models.json"); + const models = JSON.parse(readFileSync(modelsPath, "utf8")); + const declared: Record = { temperature: 0.2, top_k: 40, stop: [""] }; + models.providers.mock.models[0].samplingParams = declared; + writeFileSync(modelsPath, JSON.stringify(models)); + const result = await runPi(sandbox, "Explore this project and summarize it", { env: langfuseEnv(capture) }); + assert.equal(result.status, 0, `pi failed: ${result.stderr}`); + await waitForRequests(capture, 1); + + const sent = wired.payloads(); + const params = generationParameters(capture.spans()); + assert.equal(params.length, sent.length, "one generation per provider request"); + const keys = Object.keys(declared); + for (const [i, p] of params.entries()) { + assert.equal(p?.max_tokens, 8192, `generation ${i + 1} must still carry the cap`); + assert.deepEqual( + keys.filter((k) => p?.[k] !== undefined), + keys.filter((k) => sent[i]?.[k] !== undefined), + `generation ${i + 1} must report exactly the declared parameters the wire carried`, + ); + assert.equal(p?.temperature, sent[i]?.temperature, `generation ${i + 1} temperature`); + assert.equal(p?.top_k, sent[i]?.top_k, `generation ${i + 1} top_k`); + } + } finally { + wired.close(); + capture.close(); + } + }); + + it("puts the context window on every generation", async () => { + const capture = await startCaptureServer(); + try { + const sandbox = createSandbox(mock.port, { contextWindow: 18000 }); + const result = await runPi(sandbox, "Explore this project and summarize it", { env: langfuseEnv(capture) }); + assert.equal(result.status, 0, `pi failed: ${result.stderr}`); + await waitForRequests(capture, 1); + + const generations = capture.spans().filter((s) => s.name === "LLM Call"); + assert.equal(generations.length, 3); + for (const g of generations) { + assert.equal(g.attrs["langfuse.observation.metadata.provider"], "mock"); + assert.equal(Number(g.attrs["langfuse.observation.metadata.context_window"]), 18000); + } + } finally { + capture.close(); + } + }); });