diff --git a/.changeset/light-plants-return.md b/.changeset/light-plants-return.md new file mode 100644 index 000000000..3f1a80f76 --- /dev/null +++ b/.changeset/light-plants-return.md @@ -0,0 +1,6 @@ +--- +"@braintrust/deepseek-harness": major +"braintrust": major +--- + +feat!: Add `BRAINTRUST_CAPTURE_ATTACHMENTS` env var option and `captureAttachments` logger option for all integrations diff --git a/e2e/helpers/attachment-capture-assertions.ts b/e2e/helpers/attachment-capture-assertions.ts new file mode 100644 index 000000000..122a944c6 --- /dev/null +++ b/e2e/helpers/attachment-capture-assertions.ts @@ -0,0 +1,67 @@ +import { expect, test } from "vitest"; +import { withScenarioHarness } from "./scenario-harness"; +import { findLatestSpan } from "./trace-selectors"; +import { captureCases } from "./attachment-capture-scenario.mjs"; + +export function defineAttachmentCaptureTests(options: { + scenarioDir: string; + originalScenarioDir: string; + provider: "openai" | "anthropic"; + packageName: string; + variantKey: string; +}) { + for (const mode of ["wrapped", "auto"]) { + test(`attachment capture policies (${options.variantKey}, ${mode})`, async () => { + await withScenarioHarness(async (harness) => { + await harness.runNodeScenarioDir({ + scenarioDir: options.scenarioDir, + entry: "scenario.capture.mjs", + nodeArgs: mode === "auto" ? ["--import", "braintrust/hook.mjs"] : [], + env: { + CAPTURE_PACKAGE_NAME: options.packageName, + CAPTURE_WRAPPED: String(mode === "wrapped"), + BRAINTRUST_CAPTURE_ATTACHMENTS: "", + }, + runContext: { + variantKey: options.variantKey, + originalScenarioDir: options.originalScenarioDir, + // The same Chat Completions request works on all supported SDK majors. + // Its real-provider recording lives in the existing v6 media cassette. + ...(options.provider === "openai" + ? { cassette: { variantKey: "openai-v6" } } + : {}), + }, + timeoutMs: 180_000, + }); + const events = harness.events(); + for (const testCase of captureCases) { + const root = findLatestSpan( + events, + `attachment-capture-${testCase.name}`, + ); + expect(root).toBeDefined(); + const children = events.filter( + (event) => + event.row.root_span_id === root?.row.root_span_id && + event.span.id !== root?.span.id, + ); + expect(children.length).toBeGreaterThan(0); + const payload = JSON.stringify(children); + if (testCase.enabled) + expect(payload).toContain("braintrust_attachment"); + else { + expect(payload).not.toContain("braintrust_attachment"); + expect(payload).not.toContain('"file_data":'); + expect(payload).not.toContain('"b64_json":'); + expect(payload).not.toContain('"data":'); + expect(payload).not.toContain("data:image/png;base64,"); + } + expect(payload).toContain("completion_tokens"); + expect(children.some((event) => event.row.output !== undefined)).toBe( + true, + ); + } + }); + }, 240_000); + } +} diff --git a/e2e/helpers/attachment-capture-scenario.mjs b/e2e/helpers/attachment-capture-scenario.mjs new file mode 100644 index 000000000..f8145696f --- /dev/null +++ b/e2e/helpers/attachment-capture-scenario.mjs @@ -0,0 +1,125 @@ +import assert from "node:assert/strict"; +import { readFile } from "node:fs/promises"; +import { initLogger, wrapOpenAI, wrapAnthropic } from "braintrust"; +import { MINIMAL_PNG_BASE64 } from "./media-fixtures.mjs"; +import { scopedName } from "./provider-runtime.mjs"; + +export const captureCases = [ + { name: "default", env: "", enabled: false }, + { name: "local-default", env: "", global: true, local: true, enabled: false }, + { + name: "local-enabled", + env: "false", + global: false, + local: true, + option: true, + enabled: true, + }, + { + name: "local-disabled", + env: "true", + global: true, + local: true, + option: false, + enabled: false, + }, + { + name: "local-environment", + env: "true", + global: false, + local: true, + enabled: true, + }, + { name: "global-enabled", env: "false", option: true, enabled: true }, + { name: "global-disabled", env: "true", option: false, enabled: false }, +]; + +export async function runAttachmentCaptureScenario(provider, scenarioUrl, sdk) { + const wrapped = process.env.CAPTURE_WRAPPED === "true"; + let request; + if (provider === "anthropic") { + const raw = new sdk.default({ + apiKey: process.env.ANTHROPIC_API_KEY, + baseURL: process.env.ANTHROPIC_BASE_URL, + }); + const client = wrapped ? wrapAnthropic(raw) : raw; + const data = ( + await readFile(new URL("./test-image.png", scenarioUrl)) + ).toString("base64"); + request = async () => { + const result = await client.messages.create({ + model: "claude-haiku-4-5", + max_tokens: 32, + temperature: 0, + messages: [ + { + role: "user", + content: [ + { + type: "text", + text: "Describe the attached image in one short sentence.", + }, + { + type: "image", + source: { type: "base64", media_type: "image/png", data }, + }, + ], + }, + ], + }); + assert.ok(result.content.length > 0); + }; + } else if (provider === "openai") { + const raw = new sdk.default({ + apiKey: process.env.OPENAI_API_KEY, + baseURL: process.env.OPENAI_BASE_URL, + }); + const client = wrapped ? wrapOpenAI(raw) : raw; + request = async () => { + const result = await client.chat.completions.create({ + model: "gpt-4o-mini-2024-07-18", + messages: [ + { + role: "user", + content: [ + { + type: "text", + text: "Describe this image in three words or fewer.", + }, + { + type: "image_url", + image_url: { + url: `data:image/png;base64,${MINIMAL_PNG_BASE64}`, + }, + }, + ], + }, + ], + max_tokens: 24, + temperature: 0, + }); + assert.ok(result.choices[0].message.content); + }; + } + const projectName = scopedName("tmp-luca-attachment-capture"); + for (const testCase of captureCases) { + process.env.BRAINTRUST_CAPTURE_ATTACHMENTS = testCase.env; + if (testCase.local) + initLogger({ projectName, captureAttachments: testCase.global }); + const logger = initLogger({ + projectName, + setCurrent: !testCase.local, + captureAttachments: testCase.option, + }); + await logger.traced(request, { + name: `attachment-capture-${testCase.name}`, + event: { + metadata: { + scenario: `${provider}-instrumentation`, + testRunId: process.env.BRAINTRUST_E2E_RUN_ID, + }, + }, + }); + await logger.flush(); + } +} diff --git a/e2e/scenarios/ai-sdk-instrumentation/scenario.test.ts b/e2e/scenarios/ai-sdk-instrumentation/scenario.test.ts index 6fd7ff261..616ea11d3 100644 --- a/e2e/scenarios/ai-sdk-instrumentation/scenario.test.ts +++ b/e2e/scenarios/ai-sdk-instrumentation/scenario.test.ts @@ -55,6 +55,7 @@ describeVariants("variants", () => { await runScenarioDir({ entry: scenario.wrapperEntry, env: { + BRAINTRUST_CAPTURE_ATTACHMENTS: "true", AI_SDK_PACKAGE_NAME: scenario.packageName, AI_SDK_OPENAI_PACKAGE_NAME: scenario.openaiModuleName, AI_SDK_SUPPORTS_EVALUATE_STRING_MODEL: String( @@ -116,6 +117,7 @@ describeVariants("variants", () => { await runNodeScenarioDir({ entry: scenario.autoEntry, env: { + BRAINTRUST_CAPTURE_ATTACHMENTS: "true", AI_SDK_PACKAGE_NAME: scenario.packageName, AI_SDK_OPENAI_PACKAGE_NAME: scenario.openaiModuleName, AI_SDK_SUPPORTS_EVALUATE_STRING_MODEL: String( diff --git a/e2e/scenarios/anthropic-instrumentation/scenario.capture.mjs b/e2e/scenarios/anthropic-instrumentation/scenario.capture.mjs new file mode 100644 index 000000000..44a743716 --- /dev/null +++ b/e2e/scenarios/anthropic-instrumentation/scenario.capture.mjs @@ -0,0 +1,10 @@ +import { runMain } from "../../helpers/provider-runtime.mjs"; +import { runAttachmentCaptureScenario } from "../../helpers/attachment-capture-scenario.mjs"; + +runMain(async () => + runAttachmentCaptureScenario( + "anthropic", + import.meta.url, + await import(process.env.CAPTURE_PACKAGE_NAME), + ), +); diff --git a/e2e/scenarios/anthropic-instrumentation/scenario.test.ts b/e2e/scenarios/anthropic-instrumentation/scenario.test.ts index f2c7e3cf6..32dc02950 100644 --- a/e2e/scenarios/anthropic-instrumentation/scenario.test.ts +++ b/e2e/scenarios/anthropic-instrumentation/scenario.test.ts @@ -1,3 +1,4 @@ +import { defineAttachmentCaptureTests } from "../../helpers/attachment-capture-assertions"; import { describe } from "vitest"; import { prepareScenarioDir, @@ -52,7 +53,10 @@ describe.concurrent("variants", () => { runScenario: async ({ runScenarioDir }) => { await runScenarioDir({ entry: scenario.wrapperEntry, - env: { ANTHROPIC_PACKAGE_NAME: scenario.dependencyName }, + env: { + BRAINTRUST_CAPTURE_ATTACHMENTS: "true", + ANTHROPIC_PACKAGE_NAME: scenario.dependencyName, + }, runContext: { variantKey: scenario.snapshotName, originalScenarioDir, @@ -76,7 +80,10 @@ describe.concurrent("variants", () => { runScenario: async ({ runNodeScenarioDir }) => { await runNodeScenarioDir({ entry: scenario.autoEntry, - env: { ANTHROPIC_PACKAGE_NAME: scenario.dependencyName }, + env: { + BRAINTRUST_CAPTURE_ATTACHMENTS: "true", + ANTHROPIC_PACKAGE_NAME: scenario.dependencyName, + }, nodeArgs: ["--import", "braintrust/hook.mjs"], runContext: { variantKey: scenario.snapshotName, @@ -98,3 +105,13 @@ describe.concurrent("variants", () => { }); } }); + +for (const scenario of anthropicScenarios) { + defineAttachmentCaptureTests({ + scenarioDir, + originalScenarioDir, + provider: "anthropic", + packageName: scenario.dependencyName, + variantKey: scenario.snapshotName, + }); +} diff --git a/e2e/scenarios/elevenlabs-instrumentation/scenario.test.ts b/e2e/scenarios/elevenlabs-instrumentation/scenario.test.ts index e4d82876d..4dbecef9f 100644 --- a/e2e/scenarios/elevenlabs-instrumentation/scenario.test.ts +++ b/e2e/scenarios/elevenlabs-instrumentation/scenario.test.ts @@ -30,6 +30,7 @@ for (const variant of ["elevenlabs-v2", "elevenlabs-v2-latest"]) { ? ["--import", "braintrust/hook.mjs"] : [], env: { + BRAINTRUST_CAPTURE_ATTACHMENTS: "true", ELEVENLABS_PACKAGE_NAME: variant.replace( "elevenlabs-", "elevenlabs-sdk-", diff --git a/e2e/scenarios/google-genai-instrumentation/scenario.test.ts b/e2e/scenarios/google-genai-instrumentation/scenario.test.ts index 7863a372f..07bff02e8 100644 --- a/e2e/scenarios/google-genai-instrumentation/scenario.test.ts +++ b/e2e/scenarios/google-genai-instrumentation/scenario.test.ts @@ -56,7 +56,10 @@ describe.concurrent("variants", () => { runScenario: async ({ runScenarioDir }) => { await runScenarioDir({ entry: scenario.wrapperEntry, - env: { GOOGLE_GENAI_PACKAGE_NAME: scenario.dependencyName }, + env: { + GOOGLE_GENAI_PACKAGE_NAME: scenario.dependencyName, + BRAINTRUST_CAPTURE_ATTACHMENTS: "true", + }, runContext: { variantKey: scenario.snapshotName, originalScenarioDir, @@ -75,7 +78,10 @@ describe.concurrent("variants", () => { runScenario: async ({ runNodeScenarioDir }) => { await runNodeScenarioDir({ entry: scenario.autoEntry, - env: { GOOGLE_GENAI_PACKAGE_NAME: scenario.dependencyName }, + env: { + GOOGLE_GENAI_PACKAGE_NAME: scenario.dependencyName, + BRAINTRUST_CAPTURE_ATTACHMENTS: "true", + }, nodeArgs: ["--import", "braintrust/hook.mjs"], runContext: { variantKey: scenario.snapshotName, diff --git a/e2e/scenarios/google-generative-ai-instrumentation/scenario.test.ts b/e2e/scenarios/google-generative-ai-instrumentation/scenario.test.ts index 043dae2f4..3d24ba178 100644 --- a/e2e/scenarios/google-generative-ai-instrumentation/scenario.test.ts +++ b/e2e/scenarios/google-generative-ai-instrumentation/scenario.test.ts @@ -30,7 +30,10 @@ describe.concurrent("variants", () => { runScenario: async ({ runScenarioDir }) => { await runScenarioDir({ entry: "scenario.ts", - env: { GOOGLE_GENERATIVE_AI_PACKAGE_NAME: scenario.dependencyName }, + env: { + GOOGLE_GENERATIVE_AI_PACKAGE_NAME: scenario.dependencyName, + BRAINTRUST_CAPTURE_ATTACHMENTS: "true", + }, runContext: { originalScenarioDir, variantKey: scenario.snapshotName, @@ -50,7 +53,10 @@ describe.concurrent("variants", () => { await runNodeScenarioDir({ entry: "scenario.mjs", nodeArgs: ["--import", "braintrust/hook.mjs"], - env: { GOOGLE_GENERATIVE_AI_PACKAGE_NAME: scenario.dependencyName }, + env: { + GOOGLE_GENERATIVE_AI_PACKAGE_NAME: scenario.dependencyName, + BRAINTRUST_CAPTURE_ATTACHMENTS: "true", + }, runContext: { originalScenarioDir, variantKey: scenario.snapshotName, diff --git a/e2e/scenarios/groq-instrumentation/scenario.test.ts b/e2e/scenarios/groq-instrumentation/scenario.test.ts index be093c3d8..10641bbbc 100644 --- a/e2e/scenarios/groq-instrumentation/scenario.test.ts +++ b/e2e/scenarios/groq-instrumentation/scenario.test.ts @@ -44,6 +44,7 @@ describe.concurrent("variants", () => { await runScenarioDir({ entry: "scenario.ts", env: { + BRAINTRUST_CAPTURE_ATTACHMENTS: "true", GROQ_AUDIO_FILE: audioFile, GROQ_PACKAGE_NAME: scenario.dependencyName, }, @@ -66,6 +67,7 @@ describe.concurrent("variants", () => { await runNodeScenarioDir({ entry: "scenario.mjs", env: { + BRAINTRUST_CAPTURE_ATTACHMENTS: "true", GROQ_AUDIO_FILE: audioFile, GROQ_PACKAGE_NAME: scenario.dependencyName, }, diff --git a/e2e/scenarios/openai-instrumentation/scenario.capture.mjs b/e2e/scenarios/openai-instrumentation/scenario.capture.mjs new file mode 100644 index 000000000..780f3f1e5 --- /dev/null +++ b/e2e/scenarios/openai-instrumentation/scenario.capture.mjs @@ -0,0 +1,10 @@ +import { runMain } from "../../helpers/provider-runtime.mjs"; +import { runAttachmentCaptureScenario } from "../../helpers/attachment-capture-scenario.mjs"; + +runMain(async () => + runAttachmentCaptureScenario( + "openai", + import.meta.url, + await import(process.env.CAPTURE_PACKAGE_NAME), + ), +); diff --git a/e2e/scenarios/openai-instrumentation/scenario.test.ts b/e2e/scenarios/openai-instrumentation/scenario.test.ts index 951ef3379..dc1fa0e69 100644 --- a/e2e/scenarios/openai-instrumentation/scenario.test.ts +++ b/e2e/scenarios/openai-instrumentation/scenario.test.ts @@ -1,3 +1,4 @@ +import { defineAttachmentCaptureTests } from "../../helpers/attachment-capture-assertions"; import { describe, it } from "vitest"; import { prepareScenarioDir, @@ -136,3 +137,13 @@ describe.concurrent("variants", () => { }); } }); + +for (const scenario of openaiScenarios) { + defineAttachmentCaptureTests({ + scenarioDir, + originalScenarioDir, + provider: "openai", + packageName: scenario.dependencyName, + variantKey: scenario.snapshotName, + }); +} diff --git a/e2e/scenarios/strands-agent-sdk-instrumentation/scenario.test.ts b/e2e/scenarios/strands-agent-sdk-instrumentation/scenario.test.ts index 6d28b2730..cbcfd1167 100644 --- a/e2e/scenarios/strands-agent-sdk-instrumentation/scenario.test.ts +++ b/e2e/scenarios/strands-agent-sdk-instrumentation/scenario.test.ts @@ -49,7 +49,10 @@ describe.concurrent("variants", () => { runScenario: async ({ runScenarioDir }) => { await runScenarioDir({ entry: scenario.wrapperEntry, - env: { STRANDS_AGENT_SDK_PACKAGE_NAME: scenario.dependencyName }, + env: { + BRAINTRUST_CAPTURE_ATTACHMENTS: "true", + STRANDS_AGENT_SDK_PACKAGE_NAME: scenario.dependencyName, + }, runContext: { variantKey: scenario.variantKey, originalScenarioDir, @@ -68,7 +71,10 @@ describe.concurrent("variants", () => { runScenario: async ({ runNodeScenarioDir }) => { await runNodeScenarioDir({ entry: scenario.autoEntry, - env: { STRANDS_AGENT_SDK_PACKAGE_NAME: scenario.dependencyName }, + env: { + BRAINTRUST_CAPTURE_ATTACHMENTS: "true", + STRANDS_AGENT_SDK_PACKAGE_NAME: scenario.dependencyName, + }, nodeArgs: ["--import", "braintrust/hook.mjs"], runContext: { variantKey: scenario.variantKey, diff --git a/integrations/deepseek-harness/README.md b/integrations/deepseek-harness/README.md index 339c40783..24f89e2e2 100644 --- a/integrations/deepseek-harness/README.md +++ b/integrations/deepseek-harness/README.md @@ -30,8 +30,10 @@ To override the defaults, add the following entry to the profile's ``` The optional configuration fields are `apiKey`, `projectName`, `metadata`, -`orgName`, and `appUrl`. `metadata` is added to every top-level user-turn trace, -and `projectName` defaults to `DeepSeek Harness`. You can set `apiKey` -under **Settings > Plugins** in Harness; it is declared as a secret field so -the settings UI handles it as a credential. A configured `apiKey` takes -precedence over the `BRAINTRUST_API_KEY` environment variable. +`orgName`, `appUrl`, and `captureAttachments`. `metadata` is added to every +top-level user-turn trace, and `projectName` defaults to `DeepSeek Harness`. +You can set `apiKey` under **Settings > Plugins** in Harness; it is declared as +a secret field so the settings UI handles it as a credential. A configured +`apiKey` takes precedence over the `BRAINTRUST_API_KEY` environment variable. +Set `captureAttachments: true` to upload Harness images as attachments; when it +is omitted, the `BRAINTRUST_CAPTURE_ATTACHMENTS` environment variable decides. diff --git a/integrations/deepseek-harness/src/index.test.ts b/integrations/deepseek-harness/src/index.test.ts index 6d2c6dcf8..ca318eee3 100644 --- a/integrations/deepseek-harness/src/index.test.ts +++ b/integrations/deepseek-harness/src/index.test.ts @@ -49,10 +49,12 @@ const mock = vi.hoisted(() => { flush: vi.fn(async () => undefined), }; const initLogger = vi.fn(() => logger); + const captureAttachmentsEnabled = vi.fn(() => true); return { Attachment: MockAttachment, attachments, + captureAttachmentsEnabled, initLogger, logger, reset() { @@ -61,12 +63,14 @@ const mock = vi.hoisted(() => { spans.length = 0; initLogger.mockClear(); logger.flush.mockClear(); + captureAttachmentsEnabled.mockReturnValue(true); }, spans, }; }); vi.mock("braintrust", () => ({ + _internalCaptureAttachmentsEnabled: mock.captureAttachmentsEnabled, Attachment: mock.Attachment, initLogger: mock.initLogger, NOOP_SPAN: {}, @@ -148,6 +152,14 @@ describe("DeepSeek Harness plugin", () => { expect(ConfigSchema.dict?.apiKey?.meta.role).toBe("secret"); }); + test("passes the attachment capture setting to its logger", () => { + createContext({ captureAttachments: true }); + + expect(mock.initLogger).toHaveBeenCalledWith( + expect.objectContaining({ captureAttachments: true }), + ); + }); + test("only logs the supported tool definition fields", async () => { const harness = createContext(); const stream = harness.listener("llm/stream")( @@ -228,88 +240,104 @@ describe("DeepSeek Harness plugin", () => { } }); - test("converts Harness image references to Braintrust attachments", async () => { - const harness = createContext(); - const currentSession = session(); - const image = { - attachmentId: "image-1", - mediaType: "image/png", - bytes: 3, - width: 1, - height: 1, - name: "diagram.png", - }; - const message = { - role: "user", - content: [ - { type: "text", text: "describe this" }, - { type: "image", attachment: image }, - ], - source: { kind: "user" }, - }; - - harness.listener("session/event")(currentSession, { - type: "turn/start", - data: { turn: 1 }, - }); - harness.listener("session/event")(currentSession, { - type: "user/message", - data: message, - }); - await consume( - harness.listener("llm/stream")( - { - provider: "replay", - model: "replay-model", - messages: [message], - sessionId: currentSession.id, - }, - () => - (async function* () { - yield { type: "finish", reason: { kind: "stop" } }; - })(), - ), - ); - await harness.listener("session/flush")(currentSession); - - expect(harness.readImage).toHaveBeenCalledOnce(); - expect(harness.readImage).toHaveBeenCalledWith(image); - expect(mock.attachments).toHaveLength(1); - expect(mock.attachments[0]?.params).toMatchObject({ - filename: "diagram.png", - contentType: "image/png", - }); - expect(new Uint8Array(mock.attachments[0]?.params.data ?? [])).toEqual( - new Uint8Array([1, 2, 3]), - ); + test.each([true, false])( + "processes Harness images with capture=%s", + async (captureAttachments) => { + mock.captureAttachmentsEnabled.mockReturnValue(captureAttachments); + const harness = createContext(); + const currentSession = session(); + const image = { + attachmentId: "image-1", + mediaType: "image/png", + bytes: 3, + width: 1, + height: 1, + name: "diagram.png", + }; + const message = { + role: "user", + content: [ + { type: "text", text: "describe this" }, + { type: "image", attachment: image }, + ], + source: { kind: "user" }, + }; - const expectedImagePart = { - type: "image_url", - image_url: { url: mock.attachments[0] }, - }; - const turn = mock.spans.find( - (span) => span.args.name === "deepseek_harness.turn", - ); - const model = mock.spans.find( - (span) => span.args.name === "deepseek_harness.step", - ); - expect(turn?.logs).toContainEqual({ - input: [ - { - role: "user", - content: [{ type: "text", text: "describe this" }, expectedImagePart], - }, - ], - }); - expect(model?.logs).toContainEqual({ - input: [ - { - role: "user", - content: [{ type: "text", text: "describe this" }, expectedImagePart], - }, - ], - }); - }); + harness.listener("session/event")(currentSession, { + type: "turn/start", + data: { turn: 1 }, + }); + harness.listener("session/event")(currentSession, { + type: "user/message", + data: message, + }); + await consume( + harness.listener("llm/stream")( + { + provider: "replay", + model: "replay-model", + messages: [message], + sessionId: currentSession.id, + }, + () => + (async function* () { + yield { type: "finish", reason: { kind: "stop" } }; + })(), + ), + ); + await harness.listener("session/flush")(currentSession); + + let expectedImagePart: Record; + if (captureAttachments) { + expect(harness.readImage).toHaveBeenCalledOnce(); + expect(harness.readImage).toHaveBeenCalledWith(image); + expect(mock.attachments).toHaveLength(1); + expect(mock.attachments[0]?.params).toMatchObject({ + filename: "diagram.png", + contentType: "image/png", + }); + expect(new Uint8Array(mock.attachments[0]?.params.data ?? [])).toEqual( + new Uint8Array([1, 2, 3]), + ); + expectedImagePart = { + type: "image_url", + image_url: { url: mock.attachments[0] }, + }; + } else { + expect(harness.readImage).not.toHaveBeenCalled(); + expect(mock.attachments).toHaveLength(0); + expectedImagePart = { type: "image", attachment: image }; + } + const turn = mock.spans.find( + (span) => span.args.name === "deepseek_harness.turn", + ); + const model = mock.spans.find( + (span) => span.args.name === "deepseek_harness.step", + ); + expect(turn?.logs).toContainEqual({ + input: [ + { + role: "user", + content: [ + { type: "text", text: "describe this" }, + expectedImagePart, + ], + }, + ], + }); + expect(model?.logs).toContainEqual({ + input: [ + { + role: "user", + content: [ + { type: "text", text: "describe this" }, + expectedImagePart, + ], + }, + ], + }); + }, + ); test("preserves Harness image references when attachment conversion fails", async () => { const harness = createContext(); diff --git a/integrations/deepseek-harness/src/index.ts b/integrations/deepseek-harness/src/index.ts index e790a35d8..4ce891edd 100644 --- a/integrations/deepseek-harness/src/index.ts +++ b/integrations/deepseek-harness/src/index.ts @@ -7,6 +7,7 @@ import type {} from "@deepseek-ai/dsh-session"; import type {} from "@deepseek-ai/dsh-tools"; import z from "@deepseek-ai/schemastery"; import { + _internalCaptureAttachmentsEnabled, Attachment, initLogger, NOOP_SPAN, @@ -49,6 +50,11 @@ export interface Config { orgName?: string; /** Optional Braintrust deployment URL. */ appUrl?: string; + /** + * Capture inline Harness images as attachments. + * Falls back to BRAINTRUST_CAPTURE_ATTACHMENTS when omitted. + */ + captureAttachments?: boolean; } export const Config: z = z.object({ @@ -57,6 +63,7 @@ export const Config: z = z.object({ metadata: z.dict(z.any()), orgName: z.string(), appUrl: z.string(), + captureAttachments: z.boolean(), }); type TokenMetrics = Record & { @@ -373,6 +380,7 @@ export function apply(ctx: Context, config: Config = {}): void { projectName: config.projectName ?? DEFAULT_PROJECT_NAME, ...(config.orgName ? { orgName: config.orgName } : {}), ...(config.appUrl ? { appUrl: config.appUrl } : {}), + captureAttachments: config.captureAttachments, setCurrent: false, }); const turns = new Map>(); @@ -394,6 +402,7 @@ export function apply(ctx: Context, config: Config = {}): void { const resolveImage = async ( ref: HarnessImageAttachmentRef, ): Promise => { + if (!_internalCaptureAttachmentsEnabled(logger)) return undefined; const attachmentStore = ( ctx as Context & { attachments?: HarnessAttachmentStore } ).attachments; diff --git a/js/src/attachment-capture.test.ts b/js/src/attachment-capture.test.ts new file mode 100644 index 000000000..95f09b368 --- /dev/null +++ b/js/src/attachment-capture.test.ts @@ -0,0 +1,1024 @@ +import { afterEach, beforeAll, beforeEach, expect, it, vi } from "vitest"; +import { + _internalCaptureAttachmentsEnabled, + Attachment, + BraintrustState, + CAPTURE_ATTACHMENTS, + ExternalAttachment, + startSpan, + _exportsForTestingOnly, + _internalGetGlobalState, + initLogger, +} from "./logger"; +import { configureNode } from "./node/config"; +import iso from "./isomorph"; +import { processInputAttachments } from "./wrappers/attachment-utils"; +import { processAttachmentsInInput } from "./instrumentation/providers/anthropic-instrumentation"; +import { + processAISDKCallInput, + processAISDKGenerateImageInput, + processAISDKGenerateImageOutput, + processAISDKOutput, + patchAISDKStreamingResult, +} from "./instrumentation/providers/ai-sdk-instrumentation"; +import { OpenAIAgentsTraceProcessor } from "./instrumentation/providers/openai-agents-trace-processor"; +import type { OpenAIAgentsSpan } from "./vendor-sdk-types/openai-agents"; +import { extractOllamaChatInput } from "./instrumentation/providers/ollama-instrumentation"; +import { openAIChannels } from "./instrumentation/providers/openai-channels"; + +configureNode(); +const state = _internalGetGlobalState(); +let background: ReturnType< + typeof _exportsForTestingOnly.useTestBackgroundLogger +>; +beforeAll(() => _exportsForTestingOnly.simulateLoginForTests()); +beforeEach(() => { + vi.stubEnv("BRAINTRUST_CAPTURE_ATTACHMENTS", undefined); + state.captureAttachments = undefined; + state.currentLogger = undefined; + background = _exportsForTestingOnly.useTestBackgroundLogger(); +}); +afterEach(async () => { + await background.drain(); + _exportsForTestingOnly.clearTestBackgroundLogger(); + state.captureAttachments = undefined; + state.currentLogger = undefined; + vi.restoreAllMocks(); + vi.unstubAllEnvs(); +}); +const loggerOptions = { + projectName: "tmp-luca-capture-tests", + projectId: "test-project-id", +}; +const inlineImage = { + type: "image_url", + image_url: { url: "data:image/png;base64,AQID" }, +}; + +it.each([false, true])( + "processes native OpenAI media with capture=%s", + (captureAttachments) => { + const input = [ + { + type: "input_image", + image_url: "data:image/png;base64,AQID", + detail: "low", + }, + { type: "input_image", image: "data:image/png;base64,AQID" }, + { type: "input_file", file_data: "AQID", filename: "document.pdf" }, + { type: "input_audio", input_audio: { data: "AQID", format: "wav" } }, + { type: "input_image", image_url: "https://example.com/image.png" }, + { type: "input_file", file_id: "file-123" }, + ]; + const result = processInputAttachments(input, captureAttachments); + expect(result).toEqual([ + { + type: "input_image", + detail: "low", + ...(captureAttachments ? { image_url: expect.any(Attachment) } : {}), + }, + ...(captureAttachments + ? [{ type: "input_image", image: expect.any(Attachment) }] + : []), + { + type: "input_file", + filename: "document.pdf", + ...(captureAttachments ? { file_data: expect.any(Attachment) } : {}), + }, + { + type: "input_audio", + input_audio: { + format: "wav", + ...(captureAttachments ? { data: expect.any(Attachment) } : {}), + }, + }, + ...input.slice(4), + ]); + expect(input[2].file_data).toBe("AQID"); + }, +); + +it.each([false, true])( + "applies the owning policy to AI SDK tool-result media (%s)", + async (captureAttachments) => { + initLogger({ ...loggerOptions, captureAttachments: !captureAttachments }); + const local = initLogger({ + ...loggerOptions, + setCurrent: false, + captureAttachments, + }); + const media = ["image-data", "file-data", "media"].map((type) => ({ + type, + data: "AQID", + mediaType: "image/png", + })); + const references = [ + { type: "text", text: "Screenshot captured" }, + { type: "image-url", url: "https://example.com/image.png" }, + { type: "file-id", fileId: "file-123" }, + ]; + const toolResult = { + type: "tool-result", + toolCallId: "call-1", + toolName: "screenshot", + output: { type: "content", value: [...media, ...references] }, + }; + const messages = [{ role: "tool", content: [toolResult] }]; + const decode = vi.spyOn(globalThis, "atob"); + await local.traced(() => { + const input = processAISDKCallInput({ + model: { modelId: "model", provider: "test" }, + messages, + }).input; + const output = processAISDKOutput({ response: { messages } }, []); + const expected = [ + ...media.map(({ data: _data, ...part }) => ({ + ...part, + ...(captureAttachments ? { data: expect.any(Attachment) } : {}), + })), + ...references, + ]; + expect(input).toHaveProperty( + "messages.0.content.0.output.value", + expected, + ); + expect(output).toHaveProperty( + "response.messages.0.content.0.output.value", + expected, + ); + }); + if (!captureAttachments) expect(decode).not.toHaveBeenCalled(); + expect(toolResult.output.value).toEqual([...media, ...references]); + expect(media.every((part) => part.data === "AQID")).toBe(true); + }, +); + +it.each([false, true])( + "processes AI SDK media throughout output containers with capture=%s", + (captureAttachments) => { + initLogger({ ...loggerOptions, captureAttachments }); + const readBytes = vi.fn(() => "AQID"); + const generatedFile = { + mediaType: "image/png", + get base64() { + return readBytes(); + }, + }; + const file = { type: "file", file: generatedFile }; + const message = { + role: "assistant", + content: [ + { type: "file", mediaType: "image/png", data: "AQID" }, + { type: "text", text: "Done" }, + ], + }; + const result = processAISDKOutput( + { + content: [file], + steps: [{ content: [file], response: { messages: [message] } }], + response: { + messages: [message], + timestamp: new Date("2026-09-28T00:00:00Z"), + }, + usage: { inputTokens: 1, outputTokens: 2 } as Parameters< + typeof processAISDKOutput + >[0]["usage"], + }, + [], + ); + expect(result).toHaveProperty( + "content", + captureAttachments + ? [{ type: "file", file: expect.any(Attachment) }] + : [], + ); + expect(result).toHaveProperty( + "steps.0.content", + captureAttachments + ? [{ type: "file", file: expect.any(Attachment) }] + : [], + ); + for (const path of [ + "response.messages.0.content", + "steps.0.response.messages.0.content", + ]) { + expect(result).toHaveProperty(path, [ + { + type: "file", + mediaType: "image/png", + ...(captureAttachments ? { data: expect.any(Attachment) } : {}), + }, + { type: "text", text: "Done" }, + ]); + } + expect(result).toHaveProperty("usage", { inputTokens: 1, outputTokens: 2 }); + expect(result).toHaveProperty( + "response.timestamp", + "2026-09-28T00:00:00.000Z", + ); + if (!captureAttachments) expect(readBytes).not.toHaveBeenCalled(); + expect(message.content[0].data).toBe("AQID"); + }, +); + +it.each([false, true])( + "uses the active Agents processor's owning policy (%s)", + async (captureAttachments) => { + const logger = initLogger({ + ...loggerOptions, + setCurrent: false, + captureAttachments, + }); + const processor = new OpenAIAgentsTraceProcessor({ logger }); + const trace = { + type: "trace" as const, + traceId: "media-trace", + name: "media trace", + groupId: null, + }; + const span: OpenAIAgentsSpan = { + type: "trace.span", + traceId: trace.traceId, + spanId: "media-span", + parentId: null, + startedAt: null, + endedAt: null, + error: null, + spanData: { + type: "response", + _input: [ + { type: "input_image", image: "data:image/png;base64,AQID" }, + { type: "input_file", file: "data:application/pdf;base64,AQID" }, + { type: "input_file", file: "https://example.com/report.pdf" }, + { type: "input_file", file: { id: "file-123" } }, + { + type: "input_file", + file: { url: "https://example.com/report.pdf" }, + }, + ], + _response: { + output: [{ type: "image_generation_call", result: "AQID" }], + }, + }, + }; + await processor.onTraceStart(trace); + await processor.onSpanStart(span); + initLogger({ ...loggerOptions, captureAttachments: !captureAttachments }); + await processor.onSpanEnd(span); + const rows = await background.drain(); + const output = rows.find( + (row) => "output" in row && row.output !== undefined, + ); + expect(output).toHaveProperty( + "output", + captureAttachments + ? [{ type: "image_generation_call", result: expect.any(Attachment) }] + : [], + ); + expect(output).toHaveProperty("input", [ + ...(captureAttachments + ? [ + { type: "input_image", image: expect.any(Attachment) }, + { type: "input_file", file: expect.any(Attachment) }, + ] + : []), + { type: "input_file", file: "https://example.com/report.pdf" }, + { type: "input_file", file: { id: "file-123" } }, + { type: "input_file", file: { url: "https://example.com/report.pdf" } }, + ]); + expect( + processor._traceSpans.get(trace.traceId)?.metadata.lastOutput, + ).toEqual(output && "output" in output ? output.output : undefined); + }, +); + +it.each([false, true])( + "keeps AI SDK streamed media on its owning policy (%s)", + async (captureAttachments) => { + const local = initLogger({ + ...loggerOptions, + setCurrent: false, + captureAttachments, + }); + const span = local.startSpan({ name: "streamed media" }); + const part = { + type: "file", + data: new Uint8Array([1, 2, 3]), + mediaType: "image/png", + }; + const result = { + fullStream: (async function* () { + yield part; + })(), + content: [part], + text: "Done", + }; + expect( + patchAISDKStreamingResult({ + defaultDenyOutputPaths: [], + endEvent: {}, + result, + span, + startTime: Date.now() / 1000, + }), + ).toBe(true); + initLogger({ ...loggerOptions, captureAttachments: !captureAttachments }); + for await (const chunk of result.fullStream) expect(chunk).toBe(part); + const rows = await background.drain(); + expect(rows.find((row) => "output" in row)).toHaveProperty("output", { + content: [ + { + type: "file", + mediaType: "image/png", + ...(captureAttachments ? { data: expect.any(Attachment) } : {}), + }, + ], + text: "Done", + }); + expect(part.data).toEqual(new Uint8Array([1, 2, 3])); + }, +); + +it.each([false, true])( + "applies the active Agents policy to speech and transcription (%s)", + async (captureAttachments) => { + const local = initLogger({ + ...loggerOptions, + setCurrent: false, + captureAttachments, + }); + const processor = new OpenAIAgentsTraceProcessor({ logger: local }); + const trace = { + type: "trace" as const, + traceId: "audio-trace", + name: "audio trace", + groupId: null, + }; + await processor.onTraceStart(trace); + for (const type of ["speech", "transcription"] as const) { + const audio = { data: "AQID", format: "pcm" }; + const span: OpenAIAgentsSpan = { + type: "trace.span", + traceId: trace.traceId, + spanId: type, + parentId: null, + startedAt: null, + endedAt: null, + error: null, + spanData: + type === "speech" + ? { type, input: "Hello", output: audio } + : { type, input: audio, output: "Hello" }, + }; + await processor.onSpanStart(span); + initLogger({ ...loggerOptions, captureAttachments: !captureAttachments }); + await processor.onSpanEnd(span); + const payload = (await background.drain()).find((row) => "output" in row); + expect(payload).toHaveProperty(type === "speech" ? "output" : "input", { + format: "pcm", + ...(captureAttachments ? { data: expect.any(Attachment) } : {}), + }); + expect(audio.data).toBe("AQID"); + } + }, +); + +it.each([ + [false, false], + [false, true], + [true, false], + [true, true], +])( + "applies capture=%s to AI SDK image-edit inputs and masks (URL object: %s)", + async (captureAttachments, maskAsURL) => { + initLogger({ ...loggerOptions, captureAttachments: !captureAttachments }); + const local = initLogger({ + ...loggerOptions, + setCurrent: false, + captureAttachments, + }); + const data = new URL("data:image/png;base64,AQID"); + const remote = new URL("https://example.com/image.png"); + const manual = new Attachment({ + data: "manual", + filename: "manual.png", + contentType: "image/png", + }); + const prompt = { + images: [ + "AQID", + new Uint8Array([1, 2, 3]), + data, + remote, + remote.href, + manual, + ], + mask: maskAsURL ? data : data.href, + text: "Edit these images", + }; + const decode = vi.spyOn(globalThis, "atob"); + const result = await local.traced(() => + processAISDKGenerateImageInput({ + model: { modelId: "image-model", provider: "test" }, + prompt, + }), + ); + expect(result.input).toHaveProperty( + "prompt", + captureAttachments + ? { + images: [ + expect.any(Attachment), + expect.any(Attachment), + expect.any(Attachment), + remote, + remote.href, + manual, + ], + mask: expect.any(Attachment), + text: prompt.text, + } + : { images: [remote, remote.href, manual], text: prompt.text }, + ); + if (!captureAttachments) expect(decode).not.toHaveBeenCalled(); + expect(prompt.images).toHaveLength(6); + expect(prompt.mask).toBe(maskAsURL ? data : data.href); + }, +); + +it("preserves AI SDK image-edit inputs when conversion fails", () => { + initLogger({ ...loggerOptions, captureAttachments: true }); + const prompt = { + images: ["not valid base64!"], + mask: "not valid base64!", + text: "Edit", + }; + const result = processAISDKGenerateImageInput({ + model: { modelId: "model", provider: "test" }, + prompt, + }); + expect(result.input).toHaveProperty("prompt", prompt); +}); + +it.each([false, true])( + "applies capture=%s to AI SDK content parts containing data URL objects", + (captureAttachments) => { + initLogger({ ...loggerOptions, captureAttachments }); + const data = new URL("data:image/png;base64,AQID"); + const content = [ + { type: "image", image: data }, + { type: "file", data, mediaType: "image/png" }, + { type: "image_url", image_url: { url: data } }, + ]; + const result = processAISDKCallInput({ + model: { modelId: "model", provider: "test" }, + messages: [{ role: "user", content }], + }); + expect(result.input).toHaveProperty("messages", [ + { + role: "user", + content: captureAttachments + ? [ + { type: "image", image: expect.any(Attachment) }, + { + type: "file", + data: expect.any(Attachment), + mediaType: "image/png", + }, + { type: "image_url", image_url: { url: expect.any(Attachment) } }, + ] + : [{ type: "file", mediaType: "image/png" }], + }, + ]); + }, +); + +it("keeps explicit AI SDK attachments with capture disabled", () => { + initLogger({ ...loggerOptions, captureAttachments: false }); + const external = new ExternalAttachment({ + url: "s3://bucket/image.png", + filename: "image.png", + contentType: "image/png", + }); + const content = [ + { type: "image", image: external }, + { type: "file", data: external, mediaType: "image/png" }, + ]; + const { input } = processAISDKCallInput({ + model: { modelId: "model", provider: "test" }, + messages: [{ role: "user", content }], + }); + expect(input).toHaveProperty("messages.0.content", content); +}); + +it.each([true, false])( + "explicit %s overrides an opposing environment setting", + (captureAttachments) => { + vi.stubEnv("BRAINTRUST_CAPTURE_ATTACHMENTS", String(!captureAttachments)); + const logger = initLogger({ ...loggerOptions, captureAttachments }); + expect(_internalCaptureAttachmentsEnabled(logger)).toBe(captureAttachments); + expect(_internalCaptureAttachmentsEnabled()).toBe(captureAttachments); + initLogger(loggerOptions); + expect(_internalCaptureAttachmentsEnabled()).toBe(captureAttachments); + }, +); + +it("local loggers use the environment, and never inherit or update the global override", () => { + initLogger({ ...loggerOptions, captureAttachments: true }); + const local = initLogger({ ...loggerOptions, setCurrent: false }); + expect(_internalCaptureAttachmentsEnabled(local)).toBe(false); + initLogger({ + ...loggerOptions, + setCurrent: false, + captureAttachments: false, + }); + expect(_internalCaptureAttachmentsEnabled()).toBe(true); + vi.stubEnv("BRAINTRUST_CAPTURE_ATTACHMENTS", "true"); + initLogger({ ...loggerOptions, captureAttachments: false }); + expect(_internalCaptureAttachmentsEnabled(local)).toBe(true); +}); + +it("isolates SDK state defaults", () => { + const other = new BraintrustState({}); + initLogger({ ...loggerOptions, captureAttachments: true }); + const otherLogger = initLogger({ + ...loggerOptions, + state: other, + captureAttachments: false, + }); + expect(_internalCaptureAttachmentsEnabled(otherLogger)).toBe(false); + expect(other.captureAttachments).toBe(false); + expect(_internalCaptureAttachmentsEnabled()).toBe(true); +}); + +it("keeps concurrent local traces and descendants isolated after global replacement", async () => { + const enabled = initLogger({ + ...loggerOptions, + setCurrent: false, + captureAttachments: true, + }); + const disabled = initLogger({ + ...loggerOptions, + setCurrent: false, + captureAttachments: false, + }); + await Promise.all( + [enabled, disabled].map(async (logger, index) => { + await logger.traced(async (parent) => { + await Promise.resolve(); + initLogger({ ...loggerOptions, captureAttachments: index !== 0 }); + await parent.traced(async () => { + await Promise.resolve(); + const output = processInputAttachments(inlineImage); + if (index === 0) + expect(output.image_url.url).toBeInstanceOf(Attachment); + else expect(output).toBeUndefined(); + }); + }); + }), + ); +}); + +it("does not decode inline media, read local images, or evaluate generated binary getters when disabled", () => { + initLogger({ ...loggerOptions, captureAttachments: false }); + const decode = vi.spyOn(globalThis, "atob"); + const stat = vi.spyOn(iso, "statSync"); + const binary = vi.fn(() => { + throw new Error("must not access media"); + }); + expect( + processAttachmentsInInput([ + { + type: "image", + source: { type: "base64", media_type: "image/png", data: "AQID" }, + }, + ]), + ).toEqual([ + { + type: "image", + source: { type: "base64", media_type: "image/png" }, + }, + ]); + const source = Object.defineProperty( + { type: "base64", media_type: "image/png" }, + "data", + { get: binary, enumerable: true }, + ); + expect(processAttachmentsInInput({ type: "image", source })).toEqual({ + type: "image", + source: { type: "base64", media_type: "image/png" }, + }); + const image = Object.defineProperty({ mediaType: "image/png" }, "base64", { + get: binary, + }); + expect(processAISDKGenerateImageOutput({ images: [image] }, [])).toEqual({}); + expect( + extractOllamaChatInput([ + { + model: "model", + messages: [ + { role: "user", content: "look", images: ["/tmp/private.png"] }, + ], + }, + ]).input, + ).toEqual([ + { + role: "user", + content: "look", + }, + ]); + expect(decode).not.toHaveBeenCalled(); + expect(stat).not.toHaveBeenCalled(); + expect(binary).not.toHaveBeenCalled(); +}); + +it("preserves explicit attachments and remote references with capture disabled", async () => { + const logger = initLogger({ ...loggerOptions, captureAttachments: false }); + const attachment = new Attachment({ + data: new Blob(["explicit"]), + filename: "explicit.txt", + contentType: "text/plain", + }); + logger.log({ input: { attachment } }); + const rows = (await background.drain()) as Array<{ + input?: unknown; + output?: unknown; + metrics?: Record; + }>; + expect( + rows.some( + (row) => + (row.input as { attachment?: unknown })?.attachment === attachment, + ), + ).toBe(true); + const attachments: Attachment[] = []; + const event = { input: { attachment } }; + _exportsForTestingOnly.extractAttachments(event, attachments); + expect(attachments).toEqual([attachment]); + expect( + processInputAttachments({ + ...inlineImage, + image_url: { url: "https://example.com/image.png" }, + }).image_url.url, + ).toBe("https://example.com/image.png"); +}); + +it("preserves a deferred instrumentation policy when a different logger supplies the span parent", async () => { + initLogger({ ...loggerOptions, captureAttachments: true }); + const args = { name: "deferred", [CAPTURE_ATTACHMENTS]: false }; + const deferred = startSpan(args); + const child = deferred.startSpan(); + expect(_internalCaptureAttachmentsEnabled(deferred)).toBe(false); + expect(_internalCaptureAttachmentsEnabled(child)).toBe(false); + child.end(); + deferred.end(); +}); + +it.each([false, true])( + "preserves AI SDK application JSON that resembles media (capture=%s)", + async (captureAttachments) => { + const logger = initLogger({ ...loggerOptions, captureAttachments }); + const object = { type: "file", data: "report text" }; + const toolCall = { + type: "tool-call", + toolCallId: "call", + toolName: "report", + input: object, + args: object, + }; + const toolResult = { + type: "tool-result", + toolCallId: "call", + toolName: "report", + input: object, + output: { type: "json", value: object }, + result: object, + }; + const messages = [ + { role: "assistant", content: [toolCall] }, + { role: "tool", content: [toolResult] }, + ]; + const result = { + object, + output: object, + content: [toolCall, toolResult], + toolCalls: [toolCall], + toolResults: [toolResult], + response: { messages }, + steps: [{ content: [toolCall, toolResult], providerMetadata: object }], + providerMetadata: object, + }; + expect(processAISDKOutput(result, [], captureAttachments)).toEqual(result); + expect( + processAISDKCallInput({ + model: { modelId: "test", provider: "test" }, + messages, + }).input, + ).toHaveProperty("messages", messages); + + const span = logger.startSpan({ name: "structured stream" }); + const streamResult = { + partialObjectStream: (async function* () { + yield object; + })(), + object: Promise.resolve(object), + }; + expect( + patchAISDKStreamingResult({ + defaultDenyOutputPaths: [], + endEvent: {}, + result: streamResult, + span, + startTime: Date.now() / 1000, + }), + ).toBe(true); + for await (const chunk of streamResult.partialObjectStream) + expect(chunk).toEqual(object); + const rows = await background.drain(); + expect(rows.find((row) => "output" in row)).toHaveProperty( + "output.object", + object, + ); + }, +); + +it.each([false, true])( + "captures Agents audio output using the owning policy (%s)", + async (captureAttachments) => { + const logger = initLogger({ + ...loggerOptions, + captureAttachments, + setCurrent: false, + }); + const processor = new OpenAIAgentsTraceProcessor({ logger }); + const trace = { + type: "trace" as const, + traceId: "audio-output", + name: "audio output", + groupId: null, + }; + const audio = { + data: "AQID", + transcript: "Hello", + id: "audio-1", + expires_at: 123, + }; + const output = [{ choices: [{ message: { role: "assistant", audio } }] }]; + const span: OpenAIAgentsSpan = { + type: "trace.span", + traceId: trace.traceId, + spanId: "generation", + parentId: null, + startedAt: null, + endedAt: null, + error: null, + spanData: { type: "generation", output }, + }; + await processor.onTraceStart(trace); + await processor.onSpanStart(span); + initLogger({ ...loggerOptions, captureAttachments: !captureAttachments }); + await processor.onSpanEnd(span); + const rows = await background.drain(); + const row = rows.find((row) => "output" in row && row.output !== undefined); + expect(row).toHaveProperty("output.0.choices.0.message.audio", { + transcript: "Hello", + id: "audio-1", + expires_at: 123, + ...(captureAttachments ? { data: expect.any(Attachment) } : {}), + }); + expect(audio.data).toBe("AQID"); + const protocolAudio = { + type: "audio", + audio: "AQID", + format: "mp3", + transcript: "Hello", + }; + expect(processInputAttachments(protocolAudio, captureAttachments)).toEqual({ + type: "audio", + format: "mp3", + transcript: "Hello", + ...(captureAttachments ? { audio: expect.any(Attachment) } : {}), + }); + const remoteAudio = { type: "audio", audio: { id: "file-123" } }; + expect(processInputAttachments(remoteAudio, captureAttachments)).toEqual( + remoteAudio, + ); + }, +); + +it.each([false, true])( + "preserves custom Agents JSON without mutating SDK span data (%s)", + async (captureAttachments) => { + const logger = initLogger({ ...loggerOptions, captureAttachments }); + const processor = new OpenAIAgentsTraceProcessor({ logger }); + const trace = { + type: "trace" as const, + traceId: "custom-json", + name: "custom JSON", + groupId: null, + }; + const data = Object.freeze({ + input: Object.freeze({ type: "file", data: "report text" }), + output: Object.freeze({ + type: "image_generation_call", + result: "ordinary result", + }), + }); + const span: OpenAIAgentsSpan = { + type: "trace.span", + traceId: trace.traceId, + spanId: "custom", + parentId: null, + startedAt: null, + endedAt: null, + error: null, + spanData: { type: "custom", name: "report", data }, + }; + await processor.onTraceStart(trace); + await processor.onSpanStart(span); + await processor.onSpanEnd(span); + const rows = await background.drain(); + expect( + rows.find((row) => "output" in row && row.output !== undefined), + ).toMatchObject(data); + expect(span.spanData).toHaveProperty("data", data); + }, +); + +it.each([false, true])( + "applies capture policy to Chat Completions audio (%s)", + async (captureAttachments) => { + initLogger({ ...loggerOptions, captureAttachments }); + const audio = { data: "AQID", transcript: "Hello", id: "audio-1" }; + const result = { + choices: [ + { + index: 0, + message: { role: "assistant" as const, content: null, audio }, + finish_reason: "stop", + }, + ], + }; + await openAIChannels.chatCompletionsCreate.tracePromise( + async () => result, + { + arguments: [{ model: "audio-model", messages: [] }], + }, + ); + const rows = await background.drain(); + expect(rows.find((row) => "output" in row)).toHaveProperty( + "output.0.message.audio", + { + transcript: "Hello", + id: "audio-1", + ...(captureAttachments ? { data: expect.any(Attachment) } : {}), + }, + ); + expect(audio.data).toBe("AQID"); + }, +); + +it.each([false, true])( + "handles raw base64 chat files without changing references (capture=%s)", + async (captureAttachments) => { + initLogger({ ...loggerOptions, captureAttachments }); + const messages = [ + { + role: "user", + content: [ + { + type: "file", + file: { file_data: "AQID", filename: "document.pdf" }, + }, + { type: "file", file: { file_id: "file-123" } }, + { + type: "file", + file: { file_data: "https://example.com/document.pdf" }, + }, + { type: "text", text: "Summarize this document." }, + ], + }, + ]; + await openAIChannels.chatCompletionsCreate.tracePromise( + async () => ({ choices: [] }), + { arguments: [{ model: "model", messages }] }, + ); + const rows = await background.drain(); + expect(rows.find((row) => "input" in row)).toHaveProperty("input", [ + { + role: "user", + content: [ + { + type: "file", + file: { + filename: "document.pdf", + ...(captureAttachments + ? { file_data: expect.any(Attachment) } + : {}), + }, + }, + ...messages[0].content.slice(1), + ], + }, + ]); + expect(messages[0].content[0].file?.file_data).toBe("AQID"); + if (!captureAttachments) expect(JSON.stringify(rows)).not.toContain("AQID"); + }, +); + +it.each([false, true])( + "preserves parsed OpenAI output and tool arguments that resemble media (capture=%s)", + async (captureAttachments) => { + initLogger({ ...loggerOptions, captureAttachments }); + const parsed = { type: "file", data: "report text" }; + const parsedArguments = { + type: "input_image", + image_url: "ordinary input", + }; + const choices = [ + { + index: 0, + message: { + role: "assistant", + content: null, + parsed, + tool_calls: [ + { + type: "function", + id: "call-1", + function: { + name: "report", + arguments: JSON.stringify(parsedArguments), + parsed_arguments: parsedArguments, + }, + }, + ], + }, + }, + ]; + for (const channel of [ + openAIChannels.chatCompletionsCreate, + openAIChannels.betaChatCompletionsParse, + ]) { + await channel.tracePromise(async () => ({ choices }), { + arguments: [{ model: "model", messages: [] }], + }); + const rows = await background.drain(); + expect(rows.find((row) => "output" in row)).toHaveProperty( + "output", + choices, + ); + } + expect(choices[0].message.parsed).toBe(parsed); + expect(choices[0].message.tool_calls[0].function.parsed_arguments).toBe( + parsedArguments, + ); + }, +); + +it.each([ + ["mp3", "audio/mpeg", "mp3"], + ["wav", "audio/wav", "wav"], + ["flac", "audio/flac", "flac"], + ["aac", "audio/aac", "aac"], + ["opus", "audio/ogg", "ogg"], + ["pcm16", "audio/pcm", "pcm"], +] as const)( + "uses the requested %s format for generated chat audio", + async (format, contentType, extension) => { + initLogger({ ...loggerOptions, captureAttachments: true }); + const audio = { data: "AQID", transcript: "Hello", id: "audio-1" }; + const choices = [ + { index: 0, message: { role: "assistant", content: null, audio } }, + ]; + for (const channel of [ + openAIChannels.chatCompletionsCreate, + openAIChannels.betaChatCompletionsParse, + ]) { + await channel.tracePromise(async () => ({ choices }), { + arguments: [ + { + model: "audio-model", + messages: [], + audio: { format, voice: "alloy" }, + }, + ], + }); + const rows = await background.drain(); + expect(rows.find((row) => "output" in row)).toHaveProperty( + "output.0.message.audio", + { + transcript: "Hello", + id: "audio-1", + data: expect.objectContaining({ + reference: expect.objectContaining({ + content_type: contentType, + filename: `audio.${extension}`, + }), + }), + }, + ); + } + expect(audio).toEqual({ data: "AQID", transcript: "Hello", id: "audio-1" }); + }, +); diff --git a/js/src/exports.ts b/js/src/exports.ts index 9abb729d7..2312691ed 100644 --- a/js/src/exports.ts +++ b/js/src/exports.ts @@ -25,6 +25,7 @@ export { Prompt, ReadonlyAttachment, ReadonlyExperiment, + _internalCaptureAttachmentsEnabled, currentExperiment, currentLogger, currentSpan, diff --git a/js/src/instrumentation/core/channel-tracing.test.ts b/js/src/instrumentation/core/channel-tracing.test.ts index 319d24344..bd8157e31 100644 --- a/js/src/instrumentation/core/channel-tracing.test.ts +++ b/js/src/instrumentation/core/channel-tracing.test.ts @@ -27,6 +27,11 @@ import { traceAsyncChannel, traceStreamingChannel } from "./channel-tracing"; const testChannels = defineChannels( "channel-tracing-test", { + inputContext: channel< + [Record], + { ok: true }, + Record + >({ channelName: "async.input-context", kind: "async" }), asyncBinding: channel<[Record], { ok: true }>({ channelName: "async.binding", kind: "async", @@ -117,6 +122,31 @@ describe("traceAsyncChannel current span binding", () => { }); }); + it("extracts input in the caller's context", async () => { + let inputContextSpan: Span | undefined; + let createdSpan: Span | undefined; + traceAsyncChannel(testChannels.inputContext, { + name: "channel-tracing-test", + type: "function", + extractInput: (_args, _event, span) => { + inputContextSpan = currentSpan(); + createdSpan = span; + return { input: "input", metadata: undefined }; + }, + extractOutput: (result) => result, + extractMetrics: () => ({}), + }); + + await testChannels.inputContext.tracePromise( + async () => ({ ok: true as const }), + { arguments: [{}] } as any, + ); + + // AI SDK child tracing compares the caller's span with the created span. + expect(createdSpan).toBeDefined(); + expect(inputContextSpan).not.toBe(createdSpan); + }); + it("limits channel provenance to directly instrumented spans", async () => { traceAsyncChannel(testChannels.provenance, { name: "channel-parent", diff --git a/js/src/instrumentation/core/channel-tracing.ts b/js/src/instrumentation/core/channel-tracing.ts index 0ab5ad805..0416d4d33 100644 --- a/js/src/instrumentation/core/channel-tracing.ts +++ b/js/src/instrumentation/core/channel-tracing.ts @@ -1,6 +1,6 @@ import { debugLogger } from "../../debug-logger"; import type { IsoChannelHandlers, IsoTracingChannel } from "../../isomorph"; -import { _internalGetGlobalState, startSpan } from "../../logger"; +import { _internalGetGlobalState, startSpan, withCurrent } from "../../logger"; import type { Span } from "../../logger"; import { withSpanInstrumentationName, @@ -468,9 +468,8 @@ export function traceAsyncChannel( const { span, startTime } = spanData; try { - const output = config.extractOutput( - asyncEndEvent.result, - asyncEndEvent, + const output = withCurrent(span, () => + config.extractOutput(asyncEndEvent.result, asyncEndEvent), ); const metrics = config.extractMetrics( asyncEndEvent.result, @@ -595,20 +594,25 @@ export function traceStreamingChannel( let error: string | undefined; if (config.aggregateChunks) { - const aggregated = config.aggregateChunks( - chunks, - asyncEndEvent.result, - asyncEndEvent, - startTime, + const aggregateChunks = config.aggregateChunks; + const aggregated = withCurrent(span, () => + aggregateChunks( + chunks, + asyncEndEvent.result, + asyncEndEvent, + startTime, + ), ); output = aggregated.output; metrics = aggregated.metrics; metadata = aggregated.metadata; error = aggregated.error; } else { - output = config.extractOutput( - chunks as unknown as StreamingResult, - asyncEndEvent, + output = withCurrent(span, () => + config.extractOutput( + chunks as unknown as StreamingResult, + asyncEndEvent, + ), ); metrics = config.extractMetrics( chunks as unknown as StreamingResult, @@ -705,9 +709,11 @@ export function traceStreamingChannel( } | undefined; try { - const output = config.extractOutput( - asyncEndEvent.result as StreamingResult, - asyncEndEvent, + const output = withCurrent(span, () => + config.extractOutput( + asyncEndEvent.result as StreamingResult, + asyncEndEvent, + ), ); const metrics = config.extractMetrics( asyncEndEvent.result as StreamingResult, diff --git a/js/src/instrumentation/providers/ai-sdk-instrumentation.test.ts b/js/src/instrumentation/providers/ai-sdk-instrumentation.test.ts index c37257b8e..144bb11c0 100644 --- a/js/src/instrumentation/providers/ai-sdk-instrumentation.test.ts +++ b/js/src/instrumentation/providers/ai-sdk-instrumentation.test.ts @@ -15,6 +15,9 @@ const telemetryMocks = vi.hoisted(() => ({ vi.mock("../../isomorph", () => ({ default: { newTracingChannel: vi.fn(), + getEnv: vi.fn((name: string) => + name === "BRAINTRUST_CAPTURE_ATTACHMENTS" ? "true" : undefined, + ), }, })); diff --git a/js/src/instrumentation/providers/ai-sdk-instrumentation.ts b/js/src/instrumentation/providers/ai-sdk-instrumentation.ts index d96042db6..1566e8f73 100644 --- a/js/src/instrumentation/providers/ai-sdk-instrumentation.ts +++ b/js/src/instrumentation/providers/ai-sdk-instrumentation.ts @@ -15,9 +15,13 @@ import { } from "../../../util/index"; import { getCurrentUnixTimestamp } from "../../util"; import { + _internalCaptureAttachmentsEnabled, _internalStartSpanWithInitialMerge, Attachment, + BaseAttachment, + CAPTURE_ATTACHMENTS, currentSpan, + deepCopyEvent, startSpan, type Span, withCurrent, @@ -29,6 +33,8 @@ import { import { convertDataToBlob, getExtensionFromMediaType, + omitMediaData, + processInputAttachments, } from "../../wrappers/attachment-utils"; import { normalizeAISDKLoggedOutput } from "../../wrappers/ai-sdk/normalize-logged-output"; import { serializeAISDKToolsForLogging } from "../../wrappers/ai-sdk/tool-serialization"; @@ -805,6 +811,9 @@ function subscribeToHarnessContinuation( result: AISDKResult | AsyncIterable; }; const span = { + [CAPTURE_ATTACHMENTS]: _internalCaptureAttachmentsEnabled( + typeof parent === "string" ? undefined : parent, + ), end: () => endHarnessTurn(parent), log: (update: Parameters[0]) => updateHarnessTurn( @@ -867,6 +876,7 @@ function subscribeToHarnessContinuation( output: processAISDKOutput( endEvent.result, resolveDenyOutputPaths(endEvent, defaultDenyOutputPaths), + _internalCaptureAttachmentsEnabled(span), ), }); span.end(); @@ -1457,7 +1467,9 @@ const processMessage = (message: any): any => { if (Array.isArray(message.content)) { return { ...message, - content: message.content.map(processContentPart), + content: message.content + .map(processContentPart) + .filter((part: unknown) => part !== undefined), }; } @@ -1473,14 +1485,18 @@ const processMessage = (message: any): any => { const processPromptContent = (prompt: any): any => { if (Array.isArray(prompt)) { - return prompt.map(processContentPart); + return prompt + .map(processContentPart) + .filter((part: unknown) => part !== undefined); } if (prompt.content) { if (Array.isArray(prompt.content)) { return { ...prompt, - content: prompt.content.map(processContentPart), + content: prompt.content + .map(processContentPart) + .filter((part: unknown) => part !== undefined), }; } else if (typeof prompt.content === "object") { return { @@ -1497,7 +1513,12 @@ const processContentPart = (part: any): any => { if (!part || typeof part !== "object") return part; try { + if (part.type === "tool-result") return processInputAttachments(part); + const captureAttachments = _internalCaptureAttachmentsEnabled(); + if (part.type === "image" && part.image) { + if (!captureAttachments) + return isInlineMedia(part.image) ? omitMediaData(part, "image") : part; const imageAttachment = convertImageToAttachment( part.image, part.mimeType || part.mediaType, @@ -1510,11 +1531,9 @@ const processContentPart = (part: any): any => { } } - if ( - part.type === "file" && - part.data && - (part.mimeType || part.mediaType) - ) { + if (part.type === "file" && part.data) { + if (!captureAttachments) + return isInlineMedia(part.data) ? omitMediaData(part, "data") : part; const fileAttachment = convertDataToAttachment( part.data, part.mimeType || part.mediaType, @@ -1530,6 +1549,13 @@ const processContentPart = (part: any): any => { if (part.type === "image_url" && part.image_url) { if (typeof part.image_url === "object" && part.image_url.url) { + if (!captureAttachments) + return isInlineMedia(part.image_url.url) + ? omitMediaData({ + ...part, + image_url: omitMediaData(part.image_url, "url"), + }) + : part; const imageAttachment = convertImageToAttachment(part.image_url.url); if (imageAttachment) { return { @@ -1550,11 +1576,21 @@ const processContentPart = (part: any): any => { return part; }; +// Explicit attachments and remote URLs are references, not inline media. +const isInlineMedia = (value: unknown): boolean => + !(value instanceof BaseAttachment) && + !(value instanceof URL && value.protocol !== "data:") && + !(typeof value === "string" && /^https?:/i.test(value)); + const convertImageToAttachment = ( image: any, explicitMimeType?: string, ): Attachment | null => { try { + if (image instanceof URL) { + if (image.protocol !== "data:") return null; + image = image.href; + } if (typeof image === "string" && image.startsWith("data:")) { const [mimeTypeSection, base64Data] = image.split(","); const mimeType = mimeTypeSection.match(/data:(.*?);/)?.[1]; @@ -1608,6 +1644,10 @@ const convertDataToAttachment = ( if (!mimeType) return null; try { + if (data instanceof URL) { + if (data.protocol !== "data:") return null; + data = data.href; + } let blob: Blob | null = null; if (typeof data === "string" && data.startsWith("data:")) { @@ -1658,14 +1698,19 @@ export function processAISDKGenerateImageInput( } const processedPrompt = { ...prompt }; + const captureAttachments = _internalCaptureAttachmentsEnabled(); if (Array.isArray(prompt.images)) { - processedPrompt.images = prompt.images.map( - (image) => convertImageToAttachment(image, "image/png") ?? image, - ); + processedPrompt.images = captureAttachments + ? prompt.images.map( + (image) => convertImageToAttachment(image, "image/png") ?? image, + ) + : prompt.images.filter((image) => !isInlineMedia(image)); } if (prompt.mask !== undefined) { - processedPrompt.mask = - convertImageToAttachment(prompt.mask, "image/png") ?? prompt.mask; + if (captureAttachments) + processedPrompt.mask = + convertImageToAttachment(prompt.mask, "image/png") ?? prompt.mask; + else if (isInlineMedia(prompt.mask)) delete processedPrompt.mask; } return processAISDKCallInput({ @@ -2109,7 +2154,11 @@ function sanitizeAISDKCallInputValue(value: unknown, depth = 0): unknown { return undefined; } - if (value === null || typeof value !== "object") { + if ( + value === null || + typeof value !== "object" || + value instanceof BaseAttachment + ) { return value; } @@ -2206,7 +2255,11 @@ function sanitizeAISDKMetadataValue(value: unknown, depth = 0): unknown { return "[Function]"; } - if (value === null || typeof value !== "object") { + if ( + value === null || + typeof value !== "object" || + value instanceof BaseAttachment + ) { return value; } @@ -2696,6 +2749,7 @@ function prepareAISDKChildTracing( output: processAISDKOutput( output as AISDKResult, activeEntry.denyOutputPaths, + _internalCaptureAttachmentsEnabled(span), ), metrics, ...mergeMetadataPayload(metadataPayload, missingUsageMetadata), @@ -3184,6 +3238,7 @@ export function patchAISDKStreamingResult(args: { } const resultRecord = result as Record; + const captureAttachments = _internalCaptureAttachmentsEnabled(span); attachKnownResultPromiseHandlers(resultRecord); let finalized = false; const finalize = ( @@ -3266,6 +3321,7 @@ export function patchAISDKStreamingResult(args: { const processedOutput = await processAISDKStreamingOutput( result, resolveDenyOutputPaths(endEvent, defaultDenyOutputPaths), + captureAttachments, ); const output = transformOutput ? transformOutput(processedOutput) @@ -3621,8 +3677,13 @@ function createPatchedAsyncIterable( async function processAISDKStreamingOutput( result: AISDKResult, denyOutputPaths: string[], + captureAttachments: boolean, ): Promise | AISDKResult> { - const output = processAISDKOutput(result, denyOutputPaths); + const output = processAISDKOutput( + result, + denyOutputPaths, + captureAttachments, + ); if (!output || typeof output !== "object") { return output; @@ -3817,14 +3878,48 @@ function isAsyncGenerator(value: unknown): value is AsyncGenerator { export function processAISDKOutput( output: AISDKResult, denyOutputPaths: string[], + captureAttachments = _internalCaptureAttachmentsEnabled(), ): Record | AISDKResult { if (!output) return output; const merged = extractSerializableOutputFields(output); + // Only provider content containers contain media. Structured outputs, tool + // arguments, and metadata can use the same field names for ordinary JSON. + const processMediaContainers = ( + container: Record, + ): Record => + Object.fromEntries( + Object.entries(container).map(([key, value]) => { + if ( + ["content", "messages", "responseMessages", "toolResults"].includes( + key, + ) && + Array.isArray(value) + ) { + return [key, processInputAttachments(value, captureAttachments)]; + } + if (key === "steps" && Array.isArray(value)) { + return [ + key, + value.map((step) => + isObject(step) ? processMediaContainers(step) : step, + ), + ]; + } + if (key === "response" && isObject(value)) { + return [key, processMediaContainers(value)]; + } + return [key, value]; + }), + ); const deleteOutputPaths = denyOutputPaths.filter((path) => path.toLowerCase().endsWith("headers"), ); - const sanitized = omit(merged, denyOutputPaths, deleteOutputPaths); + const sanitized = omit( + processMediaContainers(merged), + denyOutputPaths, + deleteOutputPaths, + ); // Transport payloads can contain nested provider request/response headers; keep // user/model/tool payload fields named "headers" outside these roots intact. @@ -3917,9 +4012,10 @@ export function processAISDKGenerateImageOutput( omit(summarized, denyOutputPaths), ) as Record; if (generatedFiles.length > 0) { - loggedOutput.images = generatedFiles.map((file, index) => - convertAISDKGeneratedFileToAttachment(file, index), - ); + const images = generatedFiles + .map((file, index) => convertAISDKGeneratedFileToAttachment(file, index)) + .filter((image) => image !== undefined); + if (images.length) loggedOutput.images = images; } return loggedOutput; @@ -3933,6 +4029,8 @@ function convertAISDKGeneratedFileToAttachment( return file; } + if (!_internalCaptureAttachmentsEnabled()) return undefined; + const generatedFile = file as AISDKGeneratedFile & Record; const generatedMediaType = safeSerializableFieldRead( generatedFile, @@ -4300,13 +4398,6 @@ function extractGatewayRoutingInfo(result: AISDKResult): { return null; } -/** - * Deep copy an object via JSON serialization. - */ -function deepCopy(obj: Record): Record { - return JSON.parse(JSON.stringify(obj)); -} - /** * Parse a JSON path string into an array of keys. */ @@ -4406,7 +4497,7 @@ function omit( paths: string[], deletePaths: string[] = [], ): Record { - const result = deepCopy(obj); + const result = deepCopyEvent({ output: obj }).output; const deletePathSet = new Set(deletePaths); for (const path of paths) { diff --git a/js/src/instrumentation/providers/ai-sdk-v7-telemetry.test.ts b/js/src/instrumentation/providers/ai-sdk-v7-telemetry.test.ts index 3805b2dd3..6c6bbffd7 100644 --- a/js/src/instrumentation/providers/ai-sdk-v7-telemetry.test.ts +++ b/js/src/instrumentation/providers/ai-sdk-v7-telemetry.test.ts @@ -1,5 +1,6 @@ import { afterEach, beforeAll, beforeEach, describe, expect, it } from "vitest"; import { + Attachment, _exportsForTestingOnly, currentSpan, initLogger, @@ -41,6 +42,61 @@ describe("braintrustAISDKTelemetry", () => { _exportsForTestingOnly.clearTestBackgroundLogger(); }); + it.each([false, true])( + "keeps the owning media policy in delayed telemetry callbacks (%s)", + async (captureAttachments) => { + initLogger({ + projectName: "tmp-luca-telemetry-capture", + projectId: "test-project-id", + captureAttachments, + }); + const telemetry = braintrustAISDKTelemetry(); + const callId = "delayed-media"; + telemetry.onStart({ callId, operationId: "ai.generateText" }); + initLogger({ + projectName: "tmp-luca-telemetry-capture", + projectId: "test-project-id", + captureAttachments: !captureAttachments, + }); + const messages = [ + { + role: "user", + content: [{ type: "image", image: "data:image/png;base64,AQID" }], + }, + ]; + telemetry.onLanguageModelCallStart({ callId, messages }); + const content = [{ type: "file", mediaType: "image/png", data: "AQID" }]; + telemetry.onLanguageModelCallEnd({ callId, content }); + telemetry.onEnd({ callId, operationId: "ai.generateText", content }); + const rows = await backgroundLogger.drain(); + const outputs = rows.filter( + (row) => "output" in row && row.output !== undefined, + ); + expect(outputs).toHaveLength(2); + for (const row of outputs) + expect(row).toHaveProperty("output.content", [ + { + type: "file", + mediaType: "image/png", + ...(captureAttachments ? { data: expect.any(Attachment) } : {}), + }, + ]); + const inputRow = rows.find( + (row) => + "input" in row && + row.input && + typeof row.input === "object" && + "messages" in row.input, + ); + expect(inputRow).toHaveProperty( + "input.messages.0.content", + captureAttachments + ? [{ type: "image", image: expect.any(Attachment) }] + : [], + ); + }, + ); + it.each(["generateText", "streamText"])( "captures provider tools in %s alongside local tools", async (operation) => { diff --git a/js/src/instrumentation/providers/anthropic-instrumentation.test.ts b/js/src/instrumentation/providers/anthropic-instrumentation.test.ts index 93872d308..c945c1f2e 100644 --- a/js/src/instrumentation/providers/anthropic-instrumentation.test.ts +++ b/js/src/instrumentation/providers/anthropic-instrumentation.test.ts @@ -4,6 +4,9 @@ import { describe, it, expect, vi } from "vitest"; vi.mock("../../isomorph", () => ({ default: { newTracingChannel: vi.fn(), + getEnv: vi.fn((name: string) => + name === "BRAINTRUST_CAPTURE_ATTACHMENTS" ? "true" : undefined, + ), }, })); @@ -21,22 +24,28 @@ const aggregateAnthropicStreamChunksForTest = (chunks: unknown[]) => aggregateAnthropicStreamChunks(chunks as any); // Mock startSpan from logger -vi.mock("../../logger", () => ({ - startSpan: vi.fn(() => ({ - log: vi.fn(), - end: vi.fn(), - })), - _internalGetGlobalState: vi.fn(() => undefined), - Attachment: class Attachment { - reference: any; - constructor(opts: any) { - this.reference = { - filename: opts.filename, - content_type: opts.contentType, - }; - } - }, -})); +vi.mock("../../logger", async () => { + const { default: iso } = await import("../../isomorph"); + return { + _internalCaptureAttachmentsEnabled: () => + iso.getEnv("BRAINTRUST_CAPTURE_ATTACHMENTS") === "true", + withCurrent: (_span: unknown, callback: () => unknown) => callback(), + startSpan: vi.fn(() => ({ + log: vi.fn(), + end: vi.fn(), + })), + _internalGetGlobalState: vi.fn(() => undefined), + Attachment: class Attachment { + reference: any; + constructor(opts: any) { + this.reference = { + filename: opts.filename, + content_type: opts.contentType, + }; + } + }, + }; +}); describe("parseMetricsFromUsage", () => { it("should return empty object for null usage", () => { diff --git a/js/src/instrumentation/providers/anthropic-instrumentation.ts b/js/src/instrumentation/providers/anthropic-instrumentation.ts index e02add721..0b5f3f112 100644 --- a/js/src/instrumentation/providers/anthropic-instrumentation.ts +++ b/js/src/instrumentation/providers/anthropic-instrumentation.ts @@ -1,7 +1,9 @@ +import { omitMediaData } from "../../wrappers/attachment-utils"; import { toLoggedError } from "../core"; import { traceStreamingChannel } from "../core/channel-tracing"; import { isAsyncIterable, patchStreamIfNeeded } from "../core/stream-patcher"; import { + _internalCaptureAttachmentsEnabled, Attachment, startSpan as startBaseSpan, withCurrent, @@ -1639,15 +1641,25 @@ function convertBase64ToAttachment( /** * Process input to convert base64 attachments (images, PDFs, etc.) to Attachment objects. */ -export function processAttachmentsInInput(input: unknown): unknown { +export function processAttachmentsInInput( + input: unknown, + captureAttachments = _internalCaptureAttachmentsEnabled(), +): unknown { if (Array.isArray(input)) { - return input.map(processAttachmentsInInput); + return input + .map((value) => processAttachmentsInInput(value, captureAttachments)) + .filter((value) => value !== undefined); } if (isObject(input)) { // Check for Anthropic's content blocks with base64 data // Supports both "image" and "document" types (for PDFs, etc.) if (isAnthropicBase64ContentBlock(input)) { + if (!captureAttachments) + return omitMediaData({ + ...input, + source: omitMediaData(input.source, "data"), + }); return { ...input, source: convertBase64ToAttachment(input.source, input.type), @@ -1657,7 +1669,8 @@ export function processAttachmentsInInput(input: unknown): unknown { // Recursively process nested objects const processed: Record = {}; for (const [key, value] of Object.entries(input)) { - processed[key] = processAttachmentsInInput(value); + const result = processAttachmentsInInput(value, captureAttachments); + if (result !== undefined) processed[key] = result; } return processed; } diff --git a/js/src/instrumentation/providers/elevenlabs-instrumentation.test.ts b/js/src/instrumentation/providers/elevenlabs-instrumentation.test.ts index 0334d7638..1e525ab09 100644 --- a/js/src/instrumentation/providers/elevenlabs-instrumentation.test.ts +++ b/js/src/instrumentation/providers/elevenlabs-instrumentation.test.ts @@ -21,6 +21,7 @@ describe("ElevenLabs instrumentation", () => { beforeEach(() => { logger = _exportsForTestingOnly.useTestBackgroundLogger(); initLogger({ + captureAttachments: true, projectName: "tmp-luca-elevenlabs-unit", projectId: "test-project-id", }); @@ -29,44 +30,56 @@ describe("ElevenLabs instrumentation", () => { _exportsForTestingOnly.clearTestBackgroundLogger(); }); - it("keeps promise helpers and stream identity, and copies bytes before the caller can mutate them", async () => { - const bytes = new Uint8Array([0, 1, 2, 3]); - const stream = new ReadableStream({ - start(controller) { - controller.enqueue(bytes.subarray(1, 3)); - controller.close(); - }, - }); - const promise = Object.assign(Promise.resolve(stream), { - withRawResponse: () => - Promise.resolve({ - data: stream, - rawResponse: { - headers: new Headers({ "content-type": "audio/mpeg" }), - }, - }), - }); - const result = elevenLabsChannels.convert.invoke( - () => promise, - undefined, - ["voice", request], - {}, - ); - expect(result).toBe(promise); - expect((await result.withRawResponse()).data).toBe(stream); - const reader = stream.getReader(); - expect((await reader.read()).value).toEqual(new Uint8Array([1, 2])); - bytes.fill(9); - await reader.read(); - const spans = (await logger.drain()) as Record[]; - const output = spans.find((row: any) => row.output)?.output as any; - const attachment = output.content[0].file.file_data as Attachment; - expect(attachment).toBeInstanceOf(Attachment); - expect(output.content[0].file.byte_size).toBe(2); - expect( - new Uint8Array(await (await attachment.data()).arrayBuffer()), - ).toEqual(new Uint8Array([1, 2])); - }); + it.each([false, true])( + "keeps promise helpers and stream identity with capture=%s", + async (captureAttachments) => { + initLogger({ + projectName: "tmp-luca-elevenlabs-unit", + projectId: "test-project-id", + captureAttachments, + }); + const bytes = new Uint8Array([0, 1, 2, 3]); + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(bytes.subarray(1, 3)); + controller.close(); + }, + }); + const promise = Object.assign(Promise.resolve(stream), { + withRawResponse: () => + Promise.resolve({ + data: stream, + rawResponse: { + headers: new Headers({ "content-type": "audio/mpeg" }), + }, + }), + }); + const result = elevenLabsChannels.convert.invoke( + () => promise, + undefined, + ["voice", request], + {}, + ); + expect(result).toBe(promise); + expect((await result.withRawResponse()).data).toBe(stream); + const reader = stream.getReader(); + expect((await reader.read()).value).toEqual(new Uint8Array([1, 2])); + bytes.fill(9); + await reader.read(); + const spans = (await logger.drain()) as Record[]; + const output = spans.find((row: any) => row.output)?.output as any; + if (!captureAttachments) { + expect(output.content).toEqual([]); + return; + } + const attachment = output.content[0].file.file_data as Attachment; + expect(attachment).toBeInstanceOf(Attachment); + expect(output.content[0].file.byte_size).toBe(2); + expect( + new Uint8Array(await (await attachment.data()).arrayBuffer()), + ).toEqual(new Uint8Array([1, 2])); + }, + ); it.each(["cancel", "error", "break"])( "omits partial audio on %s", diff --git a/js/src/instrumentation/providers/elevenlabs-instrumentation.ts b/js/src/instrumentation/providers/elevenlabs-instrumentation.ts index 4a0ac6774..d6a5a2941 100644 --- a/js/src/instrumentation/providers/elevenlabs-instrumentation.ts +++ b/js/src/instrumentation/providers/elevenlabs-instrumentation.ts @@ -1,6 +1,12 @@ import { isObject, SpanTypeAttribute } from "../../../util"; import { debugLogger } from "../../debug-logger"; -import { Attachment, startSpan, withCurrent, type Span } from "../../logger"; +import { + _internalCaptureAttachmentsEnabled, + Attachment, + startSpan, + withCurrent, + type Span, +} from "../../logger"; import { INSTRUMENTATION_NAMES, withSpanInstrumentationName, @@ -70,8 +76,9 @@ export function registerElevenLabsInstrumentation(): void { file instanceof Blob ? file.type || "application/octet-stream" : "application/octet-stream"; - const blob = - file instanceof Blob + const blob = !_internalCaptureAttachmentsEnabled() + ? undefined + : file instanceof Blob ? file : file instanceof Uint8Array ? new Blob([new Uint8Array(file)], { type: contentType }) @@ -121,7 +128,8 @@ export function registerElevenLabsInstrumentation(): void { }, ([request], span) => { const file = request.file; - if (!isAsyncIterable(file)) return; + if (!_internalCaptureAttachmentsEnabled(span) || !isAsyncIterable(file)) + return; const chunks: Uint8Array[] = []; const filename = isObject(file) && typeof file.path === "string" @@ -283,6 +291,7 @@ function captureSpeech( started: number, headers?: Headers, ): void { + const captureAttachments = _internalCaptureAttachmentsEnabled(span); const format = request.outputFormat ?? "mp3_44100_128"; const formatType = format.startsWith("mp3_") ? "audio/mpeg" @@ -320,8 +329,10 @@ function captureSpeech( }); } first = false; - chunks.push(new Uint8Array(chunk)); - bytes += chunk.byteLength; + if (captureAttachments) { + chunks.push(new Uint8Array(chunk)); + bytes += chunk.byteLength; + } }; const complete = () => { if (stopped) return; @@ -365,8 +376,17 @@ function captureSpeech( finish(error); }; const timestampChunk = (chunk: ElevenLabsTimestampAudio) => { - const binary = atob(chunk.audioBase64); - observe(Uint8Array.from(binary, (character) => character.charCodeAt(0))); + if (captureAttachments) { + const binary = atob(chunk.audioBase64); + observe(Uint8Array.from(binary, (character) => character.charCodeAt(0))); + } else if (first && chunk.audioBase64) { + if (method.startsWith("stream")) { + span.log({ + metrics: { time_to_first_token: Date.now() / 1000 - started }, + }); + } + first = false; + } if (chunk.alignment || chunk.normalizedAlignment) alignments.push({ alignment: chunk.alignment, diff --git a/js/src/instrumentation/providers/eve-instrumentation.test.ts b/js/src/instrumentation/providers/eve-instrumentation.test.ts index f5b77d166..3a2b743a4 100644 --- a/js/src/instrumentation/providers/eve-instrumentation.test.ts +++ b/js/src/instrumentation/providers/eve-instrumentation.test.ts @@ -62,11 +62,13 @@ describe("braintrustEveInstrumentation", () => { initLogger({ projectName: "eve-instrumentation.test.ts", projectId: "test-project-id", + captureAttachments: false, }); }); afterEach(() => { vi.restoreAllMocks(); + vi.unstubAllEnvs(); _exportsForTestingOnly.clearTestBackgroundLogger(); }); @@ -771,130 +773,139 @@ describe("braintrustEveInstrumentation", () => { expect(step?.root_span_id).toBe(parent?.root_span_id); }); - it("resumes spans and turn metrics from Eve operation state", async () => { - const firstProvider = braintrustEveInstrumentation(); - const turnContext = providerContext(); - const modelContext = providerContext(); - const scope = ATTEMPT_SCOPE; + it.each([false, true])( + "resumes spans and turn metrics from Eve operation state (legacy: %s)", + async (legacy) => { + const firstProvider = braintrustEveInstrumentation(); + const turnContext = providerContext(); + const modelContext = providerContext(); + const scope = ATTEMPT_SCOPE; - await firstProvider.events["turn.started"]( - { - idempotencyKey: "turn:session-root:turn-0", - rootSessionId: "session-root", - sequence: 0, - sessionId: "session-root", - turnId: "turn-0", - type: "turn.started", - }, - turnContext, - ); - await firstProvider.events["model.call.started"]( - { - idempotencyKey: "model:attempt-0:0", - input: { - messages: [{ content: "Resume this call", role: "user" }], + await firstProvider.events["turn.started"]( + { + idempotencyKey: "turn:session-root:turn-0", + rootSessionId: "session-root", + sequence: 0, + sessionId: "session-root", + turnId: "turn-0", + type: "turn.started", }, - model: { - modelId: "gpt-5.4-mini", - provider: "openai.responses", + turnContext, + ); + await firstProvider.events["model.call.started"]( + { + idempotencyKey: "model:attempt-0:0", + input: { + messages: [{ content: "Resume this call", role: "user" }], + }, + model: { + modelId: "gpt-5.4-mini", + provider: "openai.responses", + }, + scope, + type: "model.call.started", }, - scope, - type: "model.call.started", - }, - modelContext, - ); + modelContext, + ); - const persisted = modelContext.state.get(); - expect(persisted).toMatchObject({ - exported: expect.any(String), - rootSpanId: expect.any(String), - spanId: expect.any(String), - }); + const persisted = { + ...(modelContext.state.get() as Record), + }; + if (legacy) delete persisted.captureAttachments; + expect(persisted).toMatchObject({ + exported: expect.any(String), + rootSpanId: expect.any(String), + spanId: expect.any(String), + }); - const resumedProvider = braintrustEveInstrumentation(); - await resumedProvider.events["model.call.completed"]( - { - content: [{ text: "Resumed output", type: "text" }], - finishReason: "stop", - idempotencyKey: "model:attempt-0:0", - scope, - type: "model.call.completed", - usage: { inputTokens: 4, outputTokens: 2 }, - }, - providerContext(persisted), - ); - const secondModelContext = providerContext(); - await resumedProvider.events["model.call.started"]( - { - idempotencyKey: "model:attempt-0:1", - input: { - messages: [{ content: "Continue this call", role: "user" }], + const resumedProvider = braintrustEveInstrumentation(); + await resumedProvider.events["model.call.completed"]( + { + content: [{ text: "Resumed output", type: "text" }], + finishReason: "stop", + idempotencyKey: "model:attempt-0:0", + scope, + type: "model.call.completed", + usage: { inputTokens: 4, outputTokens: 2 }, }, - model: { - modelId: "gpt-5.4-mini", - provider: "openai.responses", + providerContext(persisted), + ); + const secondModelContext = providerContext(); + await resumedProvider.events["model.call.started"]( + { + idempotencyKey: "model:attempt-0:1", + input: { + messages: [{ content: "Continue this call", role: "user" }], + }, + model: { + modelId: "gpt-5.4-mini", + provider: "openai.responses", + }, + scope: { ...scope, stepIndex: 1 }, + type: "model.call.started", }, - scope: { ...scope, stepIndex: 1 }, - type: "model.call.started", - }, - secondModelContext, - ); - const secondPersisted = secondModelContext.state.get(); - expect(secondPersisted).toMatchObject({ - turnMetricContributions: { - "model:attempt-0:0": { + secondModelContext, + ); + const secondPersisted = { + ...(secondModelContext.state.get() as Record), + }; + if (legacy) delete secondPersisted.captureAttachments; + expect(secondPersisted).toMatchObject({ + turnMetricContributions: { + "model:attempt-0:0": { + completion_tokens: 2, + prompt_tokens: 4, + tokens: 6, + }, + }, + }); + const finalProvider = braintrustEveInstrumentation(); + await finalProvider.events["model.call.completed"]( + { + content: [{ text: "Final resumed output", type: "text" }], + finishReason: "stop", + idempotencyKey: "model:attempt-0:1", + scope: { ...scope, stepIndex: 1 }, + type: "model.call.completed", + usage: { inputTokens: 3, outputTokens: 1 }, + }, + providerContext(secondPersisted), + ); + + const spans = (await backgroundLogger.drain()) as Array< + Record + >; + const steps = spans.filter( + (span) => span.span_attributes?.name === "eve.step", + ); + expect(steps).toHaveLength(2); + expect(steps[0]).toMatchObject({ + input: [{ content: "Resume this call", role: "user" }], + metrics: { completion_tokens: 2, prompt_tokens: 4, tokens: 6, }, - }, - }); - const finalProvider = braintrustEveInstrumentation(); - await finalProvider.events["model.call.completed"]( - { - content: [{ text: "Final resumed output", type: "text" }], - finishReason: "stop", - idempotencyKey: "model:attempt-0:1", - scope: { ...scope, stepIndex: 1 }, - type: "model.call.completed", - usage: { inputTokens: 3, outputTokens: 1 }, - }, - providerContext(secondPersisted), - ); - - const spans = (await backgroundLogger.drain()) as Array< - Record - >; - const steps = spans.filter( - (span) => span.span_attributes?.name === "eve.step", - ); - expect(steps).toHaveLength(2); - expect(steps[0]).toMatchObject({ - input: [{ content: "Resume this call", role: "user" }], - metrics: { - completion_tokens: 2, - prompt_tokens: 4, - tokens: 6, - }, - output: [ - { - finish_reason: "stop", - message: { - content: "Resumed output", - role: "assistant", + output: [ + { + finish_reason: "stop", + message: { + content: "Resumed output", + role: "assistant", + }, }, - }, - ], - }); - const turn = spans.find( - (span) => span.span_attributes?.name === "eve.turn", - ); - expect(turn?.metrics).toMatchObject({ - completion_tokens: 3, - prompt_tokens: 7, - tokens: 10, - }); - }); + ], + }); + const turn = spans.find( + (span) => span.span_attributes?.name === "eve.turn", + ); + expect(turn?.metrics).toMatchObject({ + completion_tokens: 3, + prompt_tokens: 7, + tokens: 10, + }); + }, + ); it("does not double-count replayed model completions", async () => { const provider = braintrustEveInstrumentation(); @@ -1120,159 +1131,217 @@ describe("braintrustEveInstrumentation", () => { expect(JSON.stringify(step?.input)).not.toContain("secret"); }); - it("converts inline model and tool-result media to attachments", async () => { - const provider = braintrustEveInstrumentation(); - const actionContext = providerContext(); - await provider.events["turn.started"]( - { - idempotencyKey: "turn:session-root:turn-0", - rootSessionId: "session-root", - sequence: 0, - sessionId: "session-root", - turnId: "turn-0", - type: "turn.started", - }, - providerContext(), - ); - await provider.events["model.call.started"]( - { - idempotencyKey: "model:attempt-0:0", - input: { - messages: [ - { - content: [ - { - data: "AQID", - filename: "tiny.png", - mediaType: "image/png", - type: "file", - }, - { - image: "DQ4P", - mediaType: "image/png", - type: "image", - }, + it.each( + [ + { option: false, env: "true", enabled: false }, + { option: true, env: "false", enabled: true }, + { option: undefined, env: undefined, enabled: false }, + { option: undefined, env: "true", enabled: true }, + ].flatMap((policy) => [ + { ...policy, legacy: false }, + { ...policy, legacy: true }, + ]), + )( + "applies the persisted or current policy to model media and preserves action JSON ($option, $env, legacy: $legacy)", + async ({ option, env, enabled, legacy }) => { + const captureAttachments = legacy ? !enabled : enabled; + vi.stubEnv("BRAINTRUST_CAPTURE_ATTACHMENTS", env); + const logger = initLogger({ + projectName: "eve-instrumentation.test.ts", + projectId: "test-project-id", + setCurrent: false, + captureAttachments: option, + }); + const provider = braintrustEveInstrumentation(); + const actionContext = providerContext(); + const turnContext = providerContext(); + await logger.traced(() => + provider.events["turn.started"]( + { + idempotencyKey: "turn:session-root:turn-0", + rootSessionId: "session-root", + sequence: 0, + sessionId: "session-root", + turnId: "turn-0", + type: "turn.started", + }, + turnContext, + ), + ); + initLogger({ + projectName: "eve-instrumentation.test.ts", + projectId: "test-project-id", + captureAttachments: !enabled, + }); + const persistedTurn = { + ...(turnContext.state.get() as Record), + }; + if (legacy) delete persistedTurn.captureAttachments; + // Recreate the bridge from persisted state, as Eve does across workers. + const resumedProvider = braintrustEveInstrumentation(); + await resumedProvider.events["turn.started"]( + { + idempotencyKey: "turn:session-root:turn-0", + rootSessionId: "session-root", + sequence: 0, + sessionId: "session-root", + turnId: "turn-0", + type: "turn.started", + }, + providerContext(persistedTurn), + ); + const decode = vi.spyOn(globalThis, "atob"); + await resumedProvider.events["model.call.started"]( + { + idempotencyKey: "model:attempt-0:0", + input: { + messages: [ + { + content: [ + { + data: "AQID", + filename: "tiny.png", + mediaType: "image/png", + type: "file", + }, + { + image: "DQ4P", + mediaType: "image/png", + type: "image", + }, + { + data: "not valid base64!", + filename: "invalid.png", + mediaType: "image/png", + type: "file", + }, + { type: "image", image: "https://example.com/image.png" }, + ], + role: "user", + }, + { + content: [ + { + output: { + type: "content", + value: [ + { + data: "BAUG", + mediaType: "image/png", + type: "image-data", + }, + ], + }, + toolCallId: "call-image", + toolName: "render", + type: "tool-result", + }, + ], + role: "tool", + }, + ], + }, + model: { modelId: "gpt-5.4-mini", provider: "openai.responses" }, + scope: ATTEMPT_SCOPE, + type: "model.call.started", + }, + providerContext(), + ); + await provider.events["action.started"]( + { + callId: "call-render", + idempotencyKey: "action:session-root:turn-0:call-render", + input: { + data: "report text", + type: "file", + }, + name: "render", + scope: ATTEMPT_SCOPE, + type: "action.started", + }, + actionContext, + ); + await provider.events["action.completed"]( + { + idempotencyKey: "action:session-root:turn-0:call-render", + output: { + output: { + type: "content", + value: [ { - data: "not valid base64!", - filename: "invalid.png", + data: "ordinary JSON data", mediaType: "image/png", - type: "file", - }, - ], - role: "user", - }, - { - content: [ - { - output: { - type: "content", - value: [ - { - data: "BAUG", - mediaType: "image/png", - type: "image-data", - }, - ], - }, - toolCallId: "call-image", - toolName: "render", - type: "tool-result", + type: "image-data", }, ], - role: "tool", }, - ], - }, - model: { modelId: "gpt-5.4-mini", provider: "openai.responses" }, - scope: ATTEMPT_SCOPE, - type: "model.call.started", - }, - providerContext(), - ); - await provider.events["action.started"]( - { - callId: "call-render", - idempotencyKey: "action:session-root:turn-0:call-render", - input: { - data: "data:image/png;base64,BwgJ", - mediaType: "image/png", - type: "file", - }, - name: "render", - scope: ATTEMPT_SCOPE, - type: "action.started", - }, - actionContext, - ); - await provider.events["action.completed"]( - { - idempotencyKey: "action:session-root:turn-0:call-render", - output: { - output: { - type: "content", - value: [ - { - data: new Uint8Array([10, 11, 12]), - mediaType: "image/png", - type: "image-data", - }, - ], + type: "result", }, - type: "result", + scope: ATTEMPT_SCOPE, + type: "action.completed", }, - scope: ATTEMPT_SCOPE, - type: "action.completed", - }, - actionContext, - ); + actionContext, + ); - const spans = (await backgroundLogger.drain()) as Array< - Record - >; - const step = spans.find( - (span) => span.span_attributes?.name === "eve.step", - ); - const inputAttachment = step?.input?.[0]?.content?.[0]?.data; - const imageAttachment = step?.input?.[0]?.content?.[1]?.image; - const unconvertedFile = step?.input?.[0]?.content?.[2]?.data; - const outputAttachment = - step?.input?.[1]?.content?.[0]?.output?.value?.[0]?.data; - const action = spans.find( - (span) => span.span_attributes?.name === "render", - ); - const actionInputAttachment = action?.input?.data; - const actionOutputAttachment = action?.output?.value?.[0]?.data; - expect(inputAttachment).toBeInstanceOf(Attachment); - expect(inputAttachment.reference).toMatchObject({ - content_type: "image/png", - filename: "tiny.png", - type: "braintrust_attachment", - }); - expect(imageAttachment).toBeInstanceOf(Attachment); - expect(imageAttachment.reference).toMatchObject({ - content_type: "image/png", - type: "braintrust_attachment", - }); - expect(unconvertedFile).toBe("not valid base64!"); - expect(outputAttachment).toBeInstanceOf(Attachment); - expect(outputAttachment.reference).toMatchObject({ - content_type: "image/png", - filename: "attachment.png", - type: "braintrust_attachment", - }); - expect(actionInputAttachment).toBeInstanceOf(Attachment); - expect(actionInputAttachment.reference).toMatchObject({ - content_type: "image/png", - type: "braintrust_attachment", - }); - expect(actionOutputAttachment).toBeInstanceOf(Attachment); - expect(actionOutputAttachment.reference).toMatchObject({ - content_type: "image/png", - type: "braintrust_attachment", - }); - expect(JSON.stringify(step?.input)).not.toContain("AQID"); - expect(JSON.stringify(step?.input)).not.toContain("BAUG"); - }); + const spans = (await backgroundLogger.drain()) as Array< + Record + >; + const step = spans.find( + (span) => span.span_attributes?.name === "eve.step", + ); + expect(step).toMatchObject({ + root_span_id: persistedTurn.rootSpanId, + span_parents: [persistedTurn.spanId], + }); + const inputAttachment = step?.input?.[0]?.content?.[0]?.data; + const imageAttachment = step?.input?.[0]?.content?.[1]?.image; + const unconvertedFile = step?.input?.[0]?.content?.[2]?.data; + const outputAttachment = + step?.input?.[1]?.content?.[0]?.output?.value?.[0]?.data; + const action = spans.find( + (span) => span.span_attributes?.name === "render", + ); + const actionInputAttachment = action?.input?.data; + const actionOutputAttachment = action?.output?.value?.[0]?.data; + if (captureAttachments) { + expect(inputAttachment).toBeInstanceOf(Attachment); + expect(inputAttachment.reference).toMatchObject({ + content_type: "image/png", + filename: "tiny.png", + type: "braintrust_attachment", + }); + expect(imageAttachment).toBeInstanceOf(Attachment); + expect(imageAttachment.reference).toMatchObject({ + content_type: "image/png", + type: "braintrust_attachment", + }); + expect(unconvertedFile).toBe("not valid base64!"); + expect(outputAttachment).toBeInstanceOf(Attachment); + expect(outputAttachment.reference).toMatchObject({ + content_type: "image/png", + filename: "attachment.png", + type: "braintrust_attachment", + }); + } else { + expect(inputAttachment).toBeUndefined(); + expect(imageAttachment).toBeUndefined(); + expect(unconvertedFile).toBeUndefined(); + expect(outputAttachment).toBeUndefined(); + for (const payload of ["AQID", "DQ4P", "BAUG", "not valid base64!"]) { + expect(decode).not.toHaveBeenCalledWith(payload); + } + } + expect(step?.input?.[0]?.content?.[3]).toEqual({ + type: "image", + image: "https://example.com/image.png", + }); + expect(action?.input).toEqual({ type: "file", data: "report text" }); + expect(actionInputAttachment).toBe("report text"); + expect(actionOutputAttachment).toBe("ordinary JSON data"); + expect(JSON.stringify(step?.input)).not.toContain("AQID"); + expect(JSON.stringify(step?.input)).not.toContain("BAUG"); + }, + ); it("preserves arbitrary action outputs with type fields", async () => { const provider = braintrustEveInstrumentation(); diff --git a/js/src/instrumentation/providers/eve-instrumentation.ts b/js/src/instrumentation/providers/eve-instrumentation.ts index d9c673a5e..1ac8af24c 100644 --- a/js/src/instrumentation/providers/eve-instrumentation.ts +++ b/js/src/instrumentation/providers/eve-instrumentation.ts @@ -1,6 +1,9 @@ import { debugLogger } from "../../debug-logger"; import { Attachment, + BaseAttachment, + CAPTURE_ATTACHMENTS, + _internalCaptureAttachmentsEnabled, NOOP_SPAN, _internalStartSpanWithInitialMerge, currentLogger, @@ -19,7 +22,6 @@ import { getCurrentUnixTimestamp } from "../../util"; import { convertDataToBlob, getExtensionFromMediaType, - processInputAttachments, } from "../../wrappers/attachment-utils"; import type { EveInstrumentationActionCompletedEvent, @@ -43,6 +45,7 @@ import type { type EveSpan = Pick; type EveSpanReference = { + readonly captureAttachments: boolean; readonly exported: string; readonly rootSpanId: string; readonly spanId: string; @@ -56,6 +59,7 @@ type TurnState = { rootSpanId: string; rowId: string; span: EveSpan; + captureAttachments: boolean; spanId: string; }; @@ -132,6 +136,10 @@ class ResumedEveSpan implements EveSpan { constructor(private readonly reference: EveSpanReference) {} + get [CAPTURE_ATTACHMENTS](): boolean { + return this.reference.captureAttachments; + } + get rootSpanId(): string { return this.reference.rootSpanId; } @@ -222,7 +230,7 @@ class EveBridge { ...this.metadata(event.scope.sessionId), ...modelMetadata(event.model.modelId, event.model.provider), }; - const input = capturedModelInput(event.input); + const input = capturedModelInput(event.input, turn.captureAttachments); const ids = await generateEveIds("step", event.idempotencyKey); const span = await this.startEveSpan(context, { event: { @@ -231,6 +239,7 @@ class EveBridge { metadata, }, name: "eve.step", + [CAPTURE_ATTACHMENTS]: turn.captureAttachments, parentSpanIds: { rootSpanId: turn.rootSpanId, spanId: turn.spanId, @@ -320,7 +329,7 @@ class EveBridge { ): Promise { const turn = await this.turnForScope(event.scope, context); const metadata = turn.metadata; - const input = processInputAttachments(event.input); + const input = event.input; const ids = await generateEveIds( "action", event.scope.sessionId, @@ -334,6 +343,7 @@ class EveBridge { metadata, }, name: event.name, + [CAPTURE_ATTACHMENTS]: turn.captureAttachments, parentSpanIds: { rootSpanId: turn.rootSpanId, spanId: turn.spanId, @@ -353,7 +363,7 @@ class EveBridge { error: toError(event.output.error, "Eve action failed"), }); } else if (event.output.output !== undefined) { - span.log({ output: processInputAttachments(event.output.output) }); + span.log({ output: event.output.output }); } span.end(actionEndArgs(event.acceptedAtMs)); } @@ -402,6 +412,7 @@ class EveBridge { context: EveInstrumentationHandlerContext, metadata: Record, ): Promise { + const captureAttachments = _internalCaptureAttachmentsEnabled(); const ids = await generateEveIds("turn", event.sessionId, event.turnId); const parentSpanId = event.parentLineage === undefined @@ -441,6 +452,8 @@ class EveBridge { return await this.startEveSpan(context, { event: { id: ids.rowId, metadata }, name: "eve.turn", + [CAPTURE_ATTACHMENTS]: + parentTurn?.captureAttachments ?? captureAttachments, parentSpanIds: parentSpanId === undefined ? { parentSpanIds: [], rootSpanId } @@ -452,7 +465,9 @@ class EveBridge { private async startEveSpan( context: EveInstrumentationHandlerContext, - args: Parameters[0], + args: Parameters[0] & { + [CAPTURE_ATTACHMENTS]: boolean; + }, ): Promise { const reference = readSpanReference(context.state.get()); if (reference !== undefined) { @@ -471,6 +486,7 @@ class EveBridge { try { const reference: EveSpanReference = { + captureAttachments: _internalCaptureAttachmentsEnabled(span), exported: await span.export(), rootSpanId: span.rootSpanId, spanId: span.spanId, @@ -507,6 +523,7 @@ class EveBridge { rootSpanId: span.rootSpanId, rowId: ids.rowId, span, + captureAttachments: _internalCaptureAttachmentsEnabled(span), spanId: span.spanId, }; this.turnsByKey.set(key, turn); @@ -559,6 +576,7 @@ class EveBridge { rootSpanId: span.rootSpanId, rowId: ids.rowId, span, + captureAttachments: _internalCaptureAttachmentsEnabled(span), spanId: span.spanId, }; } @@ -578,9 +596,8 @@ class EveBridge { } const ids = await generateEveIds("turn", scope.sessionId, scope.turnId); - const persistedRootSpanId = readSpanReference( - context.state.get(), - )?.rootSpanId; + const reference = readSpanReference(context.state.get()); + const persistedRootSpanId = reference?.rootSpanId; const rememberedRootSpanId = eveRootSpanIdsByTurnKey.get(key); return { key, @@ -594,6 +611,8 @@ class EveBridge { scope.turnId, )), rowId: ids.rowId, + captureAttachments: + reference?.captureAttachments ?? _internalCaptureAttachmentsEnabled(), span: NOOP_SPAN, spanId: ids.spanId, }; @@ -677,6 +696,7 @@ class EveBridge { metadata: this.metadata(event.scope.sessionId), }, name: "eve.step", + [CAPTURE_ATTACHMENTS]: turn.captureAttachments, parentSpanIds: { rootSpanId: turn.rootSpanId, spanId: turn.spanId, @@ -707,6 +727,7 @@ class EveBridge { return await this.startEveSpan(context, { event: { id: ids.rowId, metadata: turn.metadata }, name: "eve.action", + [CAPTURE_ATTACHMENTS]: turn.captureAttachments, parentSpanIds: { rootSpanId: turn.rootSpanId, spanId: turn.spanId, @@ -731,7 +752,10 @@ function readSpanReference(value: unknown): EveSpanReference | undefined { const exported = value["exported"]; const rootSpanId = value["rootSpanId"]; const spanId = value["spanId"]; + const captureAttachments = value["captureAttachments"]; if ( + (captureAttachments !== undefined && + typeof captureAttachments !== "boolean") || typeof exported !== "string" || typeof rootSpanId !== "string" || typeof spanId !== "string" @@ -739,6 +763,9 @@ function readSpanReference(value: unknown): EveSpanReference | undefined { return undefined; } return { + // State written before attachment capture was configurable has no policy. + captureAttachments: + captureAttachments ?? _internalCaptureAttachmentsEnabled(), exported, rootSpanId, spanId, @@ -754,6 +781,7 @@ function readSpanReference(value: unknown): EveSpanReference | undefined { function capturedModelInput( input: EveInstrumentationModelInput | undefined, + captureAttachments: boolean, ): unknown[] | undefined { if (input === undefined) { return undefined; @@ -763,19 +791,22 @@ function capturedModelInput( value.push({ content: input.instructions, role: "system" }); } else if (Array.isArray(input.instructions)) { for (const instruction of input.instructions) { - const projected = projectedModelMessage(instruction); + const projected = projectedModelMessage(instruction, captureAttachments); if (projected?.role === "system") { value.push(projected); } } } else if (isObject(input.instructions)) { - const instruction = projectedModelMessage(input.instructions); + const instruction = projectedModelMessage( + input.instructions, + captureAttachments, + ); if (instruction?.role === "system") { value.push(instruction); } } for (const message of input.messages) { - const projected = projectedModelMessage(message); + const projected = projectedModelMessage(message, captureAttachments); if (projected !== undefined) { value.push(projected); } @@ -785,6 +816,7 @@ function capturedModelInput( function projectedModelMessage( value: unknown, + captureAttachments: boolean, ): { content: unknown; role: string } | undefined { if (!isObject(value)) { return undefined; @@ -826,7 +858,12 @@ function projectedModelMessage( ? part["mediaType"] : "image/png"; projected.push({ - image: mediaAttachment(part["image"], mediaType, undefined), + image: mediaAttachment( + part["image"], + mediaType, + undefined, + captureAttachments, + ), ...(typeof part["mediaType"] === "string" ? { mediaType } : {}), type: "image", }); @@ -846,6 +883,7 @@ function projectedModelMessage( typeof part["filename"] === "string" ? part["filename"] : undefined, + captureAttachments, ); projected.push({ ...(data === undefined ? {} : { data }), @@ -876,7 +914,7 @@ function projectedModelMessage( typeof part["toolName"] === "string" ) { projected.push({ - output: projectedToolOutput(part["output"]), + output: projectedToolOutput(part["output"], captureAttachments), toolCallId: part["toolCallId"], toolName: part["toolName"], type: "tool-result", @@ -905,7 +943,10 @@ function projectedModelMessage( return { content: projected, role }; } -function projectedToolOutput(value: unknown): unknown { +function projectedToolOutput( + value: unknown, + captureAttachments: boolean, +): unknown { if (!isObject(value)) { return value; } @@ -945,6 +986,7 @@ function projectedToolOutput(value: unknown): unknown { part["data"], part["mediaType"], typeof part["filename"] === "string" ? part["filename"] : undefined, + captureAttachments, ); content.push({ ...(data === undefined ? {} : { data }), @@ -975,7 +1017,15 @@ function mediaAttachment( data: unknown, contentType: string, filename: string | undefined, + captureAttachments: boolean, ): unknown { + if ( + data instanceof BaseAttachment || + (data instanceof URL && data.protocol !== "data:") || + (typeof data === "string" && /^https?:/i.test(data)) + ) + return data; + if (!captureAttachments) return undefined; const blob = convertDataToBlob(data, contentType); if (blob === null) { return data; diff --git a/js/src/instrumentation/providers/google-genai-instrumentation.test.ts b/js/src/instrumentation/providers/google-genai-instrumentation.test.ts index 2d5fb2998..b495cbf65 100644 --- a/js/src/instrumentation/providers/google-genai-instrumentation.test.ts +++ b/js/src/instrumentation/providers/google-genai-instrumentation.test.ts @@ -19,6 +19,9 @@ vi.mock("../../isomorph", () => ({ }; }), newTracingChannel: vi.fn(), + getEnv: vi.fn((name: string) => + name === "BRAINTRUST_CAPTURE_ATTACHMENTS" ? "true" : undefined, + ), }, })); @@ -30,24 +33,31 @@ const mockNewTracingChannel = iso.newTracingChannel as ReturnType; const mockStartSpan = vi.mocked(startSpan); // Mock logger -vi.mock("../../logger", () => ({ - startSpan: vi.fn(() => ({ - log: vi.fn(), - end: vi.fn(), - })), - _internalGetGlobalState: vi.fn(() => undefined), - currentSpan: vi.fn(() => undefined), - withCurrent: vi.fn((_span: unknown, callback: () => unknown) => callback()), - Attachment: class MockAttachment { - reference: any; - constructor(params: any) { - this.reference = { - filename: params.filename, - content_type: params.contentType, - }; - } - }, -})); +vi.mock("../../logger", async () => { + const { default: iso } = await import("../../isomorph"); + return { + _internalCaptureAttachmentsEnabled: () => + iso.getEnv("BRAINTRUST_CAPTURE_ATTACHMENTS") === "true", + BaseAttachment: class {}, + CAPTURE_ATTACHMENTS: Symbol.for("braintrust.captureAttachments"), + startSpan: vi.fn(() => ({ + log: vi.fn(), + end: vi.fn(), + })), + _internalGetGlobalState: vi.fn(() => undefined), + currentSpan: vi.fn(() => undefined), + withCurrent: vi.fn((_span: unknown, callback: () => unknown) => callback()), + Attachment: class MockAttachment { + reference: any; + constructor(params: any) { + this.reference = { + filename: params.filename, + content_type: params.contentType, + }; + } + }, + }; +}); describe("registerGoogleGenAIInstrumentation", () => { let mockChannel: any; @@ -1088,6 +1098,132 @@ describe("registerGoogleGenAIInstrumentation", () => { }); describe("interactions.create channel subscription", () => { + it.each([ + [false, false], + [false, true], + [true, false], + [true, true], + ])( + "preserves application JSON while processing media with capture=%s, streaming=%s", + async (captureAttachments, streaming) => { + const getEnv = vi.mocked(iso.getEnv); + const originalGetEnv = getEnv.getMockImplementation()!; + getEnv.mockImplementation((name) => + name === "BRAINTRUST_CAPTURE_ATTACHMENTS" + ? String(captureAttachments) + : undefined, + ); + try { + registerGoogleGenAIInstrumentation(); + const handlers = subscribeSpy.mock.calls[3][0]; + const data = { mimeType: "application/json", data: { count: 1 } }; + const mediaShapedJSON = { + type: "image", + mime_type: "image/png", + data: "ordinary application data", + }; + const media = { type: "image", mime_type: "image/png", data: "AQID" }; + const content = [ + { + type: "function_call", + name: "report", + arguments: { data, mediaShapedJSON }, + }, + { + type: "function_result", + name: "report", + result: { data, mediaShapedJSON }, + }, + media, + ]; + const result = { + id: "interaction-json", + input: content, + steps: [ + ...content.slice(0, 2), + { type: "model_output", content: [media] }, + ], + metadata: { data, mediaShapedJSON }, + }; + const event: any = { + arguments: [ + { + model: "gemini-2.5-flash", + input: content, + agent_config: { data, mediaShapedJSON }, + }, + ], + result: streaming + ? (async function* () { + for (const [index, step] of result.steps.entries()) { + yield { event_type: "step.start", index, step }; + } + yield { + event_type: "interaction.completed", + interaction: result, + }; + })() + : result, + }; + handlers.start(event); + const span = mockStartSpan.mock.results.at(-1)!.value; + handlers.asyncEnd(event); + if (streaming) { + for await (const _chunk of event.result) { + // Consume the provider stream to trigger final trace aggregation. + } + } + const expectedContent = [ + ...content.slice(0, 2), + { + type: "image", + mime_type: "image/png", + ...(captureAttachments + ? { + data: expect.objectContaining({ + reference: { + filename: "file.png", + content_type: "image/png", + }, + }), + } + : {}), + }, + ]; + expect(span.log).toHaveBeenNthCalledWith( + 1, + expect.objectContaining({ + input: expect.objectContaining({ + input: expectedContent, + agent_config: { data, mediaShapedJSON }, + }), + metadata: expect.objectContaining({ + agent_config: { data, mediaShapedJSON }, + }), + }), + ); + const expectedSteps = [ + ...content.slice(0, 2), + { type: "model_output", content: [expectedContent[2]] }, + ]; + expect(span.log).toHaveBeenLastCalledWith( + expect.objectContaining({ + output: { + ...result, + input: expectedContent, + steps: streaming + ? expectedSteps.map((step, index) => ({ ...step, index })) + : expectedSteps, + }, + }), + ); + expect(media.data).toBe("AQID"); + } finally { + getEnv.mockImplementation(originalGetEnv); + } + }, + ); + it("subscribes to the interactions.create channel", () => { registerGoogleGenAIInstrumentation(); diff --git a/js/src/instrumentation/providers/google-genai-instrumentation.ts b/js/src/instrumentation/providers/google-genai-instrumentation.ts index 2a09d8df5..a0608b143 100644 --- a/js/src/instrumentation/providers/google-genai-instrumentation.ts +++ b/js/src/instrumentation/providers/google-genai-instrumentation.ts @@ -1,6 +1,7 @@ import { uint8ArrayToBase64 } from "../../../util/bytes"; import { getExtensionFromMediaType, + omitMediaData, processInputAttachments, } from "../../wrappers/attachment-utils"; import { debugLogger } from "../../debug-logger"; @@ -12,8 +13,10 @@ import type { } from "../core/channel-definitions"; import type { IsoChannelHandlers, IsoTracingChannel } from "../../isomorph"; import { + _internalCaptureAttachmentsEnabled, _internalGetGlobalState, Attachment, + CAPTURE_ATTACHMENTS, currentSpan, startSpan as startBaseSpan, withCurrent, @@ -67,6 +70,7 @@ type GenerateContentStreamEvent = googleGenAIInput?: Record; googleGenAIMetadata?: Record; googleGenAIStartTime?: number; + captureAttachments?: boolean; }; type SpanState = { @@ -195,7 +199,9 @@ class GoogleGenAIInstrumentationConsumer { spanState.startTime, ), ), - output: serializeGenerateContentOutput(event.result), + output: withCurrent(spanState.span, () => + serializeGenerateContentOutput(event.result), + ), }); } finally { spanState.span.end(); @@ -226,10 +232,12 @@ class GoogleGenAIInstrumentationConsumer { streamEvent.googleGenAIMetadata = extractGenerateContentMetadata(params); streamEvent.googleGenAIStartTime = getCurrentUnixTimestamp(); + streamEvent.captureAttachments = _internalCaptureAttachmentsEnabled(); }, asyncEnd: (event) => { const streamEvent = event as GenerateContentStreamEvent; patchGoogleGenAIStreamingResult({ + captureAttachments: streamEvent.captureAttachments, input: streamEvent.googleGenAIInput, metadata: streamEvent.googleGenAIMetadata, startTime: streamEvent.googleGenAIStartTime, @@ -502,7 +510,9 @@ function interceptGoogleGenAIMediaCall< try { void Promise.resolve(result).then((response) => { try { - span.log({ output: serializeOutput(response, params) }); + span.log({ + output: withCurrent(span, () => serializeOutput(response, params)), + }); } catch (error) { debugLogger.error( `Error capturing Google GenAI ${name} output:`, @@ -588,12 +598,19 @@ function logErrorAndEndSpan( } function patchGoogleGenAIStreamingResult(args: { + captureAttachments?: boolean; input: Record | undefined; metadata: Record | undefined; startTime: number | undefined; result: unknown; }): boolean { - const { input, metadata, result, startTime } = args; + const { + input, + metadata, + result, + startTime, + captureAttachments = _internalCaptureAttachmentsEnabled(), + } = args; if ( !input || @@ -618,6 +635,7 @@ function patchGoogleGenAIStreamingResult(args: { withSpanInstrumentationName( { name: "generate_content_stream", + [CAPTURE_ATTACHMENTS]: captureAttachments, spanAttributes: { type: SpanTypeAttribute.LLM, }, @@ -739,7 +757,35 @@ function patchGoogleGenAIStreamingResult(args: { if (firstTokenTime === null) { firstTokenTime = getCurrentUnixTimestamp(); } - chunks.push(nextResult.value); + chunks.push( + captureAttachments + ? nextResult.value + : { + ...nextResult.value, + candidates: nextResult.value.candidates?.map( + (candidate) => ({ + ...candidate, + content: candidate.content + ? { + ...candidate.content, + parts: candidate.content.parts?.map((part) => + part.inlineData + ? { + ...part, + // This copy is only retained for trace aggregation. + inlineData: omitMediaData( + part.inlineData, + "data", + ) as GoogleGenAIPart["inlineData"], + } + : part, + ), + } + : candidate.content, + }), + ), + }, + ); } if (nextResult.done) { @@ -748,6 +794,7 @@ function patchGoogleGenAIStreamingResult(args: { chunks, requestStartTime, firstTokenTime, + captureAttachments, ), }); } @@ -775,6 +822,7 @@ function patchGoogleGenAIStreamingResult(args: { chunks, requestStartTime, firstTokenTime, + captureAttachments, ), }); } else { @@ -850,7 +898,9 @@ function serializeGenerateContentOutput( ? { content: { ...candidate.content, - parts: candidate.content.parts.map((part) => serializePart(part)), + parts: candidate.content.parts + .map((part) => serializePart(part)) + .filter((part) => part !== undefined), }, } : {}), @@ -1231,7 +1281,7 @@ function serializeInteractionInput( ]) { const value = params[key]; if (value !== undefined) { - input[key] = serializeInteractionValue(value); + input[key] = serializeInteractionValue(value, false); } } @@ -1256,13 +1306,13 @@ function extractInteractionMetadata( ]) { const value = params[key]; if (value !== undefined) { - metadata[key] = serializeInteractionValue(value); + metadata[key] = serializeInteractionValue(value, false); } } if (Array.isArray(params.tools)) { metadata.tools = params.tools.map((tool) => - serializeInteractionValue(tool), + serializeInteractionValue(tool, false), ); } @@ -1294,7 +1344,9 @@ function serializeContentItem(item: string | GoogleGenAIContent): unknown { if (item.parts && Array.isArray(item.parts)) { return { ...item, - parts: item.parts.map((part: GoogleGenAIPart) => serializePart(part)), + parts: item.parts + .map((part: GoogleGenAIPart) => serializePart(part)) + .filter((part) => part !== undefined), }; } return item; @@ -1310,14 +1362,27 @@ function serializeContentItem(item: string | GoogleGenAIContent): unknown { /** * Serialize a part, converting inline data to Attachments. */ -function serializePart(part: GoogleGenAIPart): unknown { +function serializePart( + part: GoogleGenAIPart, + captureAttachments = _internalCaptureAttachmentsEnabled(), +): unknown { if (!part || typeof part !== "object") { return part; } + if (part.inlineData && !captureAttachments) + return omitMediaData({ + ...part, + inlineData: omitMediaData(part.inlineData, "data"), + }); if (part.inlineData && part.inlineData.data) { const { data, mimeType } = part.inlineData; - const attachment = createAttachmentFromInlineData(data, mimeType); + const attachment = createAttachmentFromInlineData( + data, + mimeType, + undefined, + captureAttachments, + ); if (attachment) { return mimeType.startsWith("image/") @@ -1510,6 +1575,7 @@ function serializeVideoInteractionOutput( function serializeInteractionValue( value: unknown, + processMedia = true, seen = new WeakSet(), ): unknown { if (value === null || value === undefined || typeof value !== "object") { @@ -1517,7 +1583,9 @@ function serializeInteractionValue( } if (Array.isArray(value)) { - return value.map((item) => serializeInteractionValue(item, seen)); + return value.map((item) => + serializeInteractionValue(item, processMedia, seen), + ); } const dict: unknown = tryToDict(value); @@ -1526,7 +1594,9 @@ function serializeInteractionValue( } if (Array.isArray(dict)) { - return dict.map((item) => serializeInteractionValue(item, seen)); + return dict.map((item) => + serializeInteractionValue(item, processMedia, seen), + ); } if (seen.has(dict)) { @@ -1536,22 +1606,48 @@ function serializeInteractionValue( seen.add(dict); try { const serialized: Record = {}; + const isMedia = + processMedia && + "type" in dict && + typeof dict.type === "string" && + ["image", "audio", "video", "document"].includes(dict.type); const mimeType = "mime_type" in dict && typeof dict.mime_type === "string" ? dict.mime_type : "mimeType" in dict && typeof dict.mimeType === "string" ? dict.mimeType : undefined; + const captureAttachments = isMedia && _internalCaptureAttachmentsEnabled(); const attachment = - mimeType && "data" in dict && dict.data !== undefined + captureAttachments && + mimeType && + "data" in dict && + dict.data !== undefined ? createAttachmentFromInlineData(dict.data, mimeType) : null; - for (const [key, entry] of Object.entries(dict)) { + for (const key of Object.keys(dict)) { + if (key === "data" && isMedia && !captureAttachments) continue; + const entry: unknown = Reflect.get(dict, key); if (key === "data" && attachment) { serialized[key] = attachment; } else { - serialized[key] = serializeInteractionValue(entry, seen); + // Only descend into provider content containers for media conversion. + // Tool arguments/results and metadata remain ordinary application JSON. + serialized[key] = serializeInteractionValue( + entry, + processMedia && + [ + "input", + "content", + "steps", + "outputs", + "output_image", + "output_audio", + "output_video", + ].includes(key), + seen, + ); } } @@ -1565,7 +1661,9 @@ function createAttachmentFromInlineData( data: unknown, mimeType?: string, filename = `file.${mimeType ? getExtensionFromMediaType(mimeType) : "bin"}`, + captureAttachments = _internalCaptureAttachmentsEnabled(), ): Attachment | null { + if (!captureAttachments) return null; if ( !( data instanceof Uint8Array || @@ -1916,6 +2014,7 @@ function aggregateGenerateContentChunks( chunks: GoogleGenAIGenerateContentResponse[], startTime: number, firstTokenTime: number | null, + captureAttachments: boolean, ): { aggregated: Record; metrics: Record; @@ -1971,7 +2070,9 @@ function aggregateGenerateContentChunks( } else if (part.executableCode) { otherParts.push({ executableCode: part.executableCode }); } else if (part.inlineData || part.fileData) { - const serializedPart = tryToDict(serializePart(part)); + const serializedPart = tryToDict( + serializePart(part, captureAttachments), + ); if (serializedPart) { otherParts.push(serializedPart); } @@ -2120,7 +2221,7 @@ function aggregateInteractionEvents( output.output_text = outputText; } if (latestUsage) { - output.usage = serializeInteractionValue(latestUsage); + output.usage = serializeInteractionValue(latestUsage, false); } const compactSteps = Array.from(steps.values()).sort( @@ -2182,7 +2283,10 @@ function compactInteractionStep(step: unknown): Record { "is_error", ]) { if (stepDict[key] !== undefined) { - compact[key] = serializeInteractionValue(stepDict[key]); + compact[key] = serializeInteractionValue( + stepDict[key], + key === "content", + ); } } diff --git a/js/src/instrumentation/providers/google-generative-ai-instrumentation.test.ts b/js/src/instrumentation/providers/google-generative-ai-instrumentation.test.ts index b73680cc3..4b669aca1 100644 --- a/js/src/instrumentation/providers/google-generative-ai-instrumentation.test.ts +++ b/js/src/instrumentation/providers/google-generative-ai-instrumentation.test.ts @@ -22,6 +22,7 @@ describe("Google Generative AI instrumentation", () => { beforeEach(() => { logger = _exportsForTestingOnly.useTestBackgroundLogger(); initLogger({ + captureAttachments: true, projectName: "tmp-luca-google-generative-ai-tests", projectId: "test-project-id", }); @@ -366,23 +367,35 @@ describe("Google Generative AI instrumentation", () => { ]); }); - it("converts image inputs without mutating the provider request", async () => { - const inlineData = { data: "aGVsbG8=", mimeType: "image/png" }; - const request = [{ inlineData }]; - await channels.generateContent.invoke( - async () => ({ response: {} }), - { model: "models/test" }, - [request], - {}, - ); - const spans = await logger.drain(); - expect(spans[0]).toHaveProperty( - "input.contents.0.parts.0.inlineData.data", - expect.any(Attachment), - ); - expect(request[0].inlineData).toBe(inlineData); - expect(inlineData.data).toBe("aGVsbG8="); - }); + it.each([false, true])( + "applies capture=%s to image inputs without mutating the provider request", + async (captureAttachments) => { + initLogger({ + projectName: "tmp-luca-google-generative-ai-tests", + projectId: "test-project-id", + captureAttachments, + }); + const inlineData = { data: "aGVsbG8=", mimeType: "image/png" }; + const request = [{ inlineData }]; + await channels.generateContent.invoke( + async () => ({ response: {} }), + { model: "models/test" }, + [request], + {}, + ); + const spans = await logger.drain(); + if (captureAttachments) { + expect(spans[0]).toHaveProperty( + "input.contents.0.parts.0.inlineData.data", + expect.any(Attachment), + ); + } else { + expect(spans[0]).toHaveProperty("input.contents.0.parts", []); + } + expect(request[0].inlineData).toBe(inlineData); + expect(inlineData.data).toBe("aGVsbG8="); + }, + ); it("keeps all original input when any attachment cannot be converted", async () => { const request = [ diff --git a/js/src/instrumentation/providers/google-generative-ai-instrumentation.ts b/js/src/instrumentation/providers/google-generative-ai-instrumentation.ts index 85e45685a..de1c6852b 100644 --- a/js/src/instrumentation/providers/google-generative-ai-instrumentation.ts +++ b/js/src/instrumentation/providers/google-generative-ai-instrumentation.ts @@ -1,3 +1,4 @@ +import { _internalCaptureAttachmentsEnabled } from "../../logger"; import { SpanTypeAttribute } from "../../../util/index"; import { debugLogger } from "../../debug-logger"; import { startSpan, withCurrent, type Span } from "../../logger"; @@ -6,7 +7,10 @@ import { withSpanInstrumentationName, } from "../../span-origin"; import { getCurrentUnixTimestamp, isObject } from "../../util"; -import { processInputAttachments } from "../../wrappers/attachment-utils"; +import { + omitMediaData, + processInputAttachments, +} from "../../wrappers/attachment-utils"; import type { GenerativeAIChat, GenerativeAIConfig, @@ -381,10 +385,13 @@ function extractInput( // Convert Google's inlineData to the shared attachment processor's file shape, // then restore the native payload shape. Any conversion failure keeps all input. function convertAttachments(value: T): T { + const captureAttachments = _internalCaptureAttachmentsEnabled(); const convert = (node: unknown): unknown => { - if (Array.isArray(node)) return node.map(convert); + if (Array.isArray(node)) + return node.map(convert).filter((item) => item !== undefined); if (!isObject(node)) return node; if (isObject(node.inlineData)) { + if (!captureAttachments) return omitMediaData(node, "inlineData"); const { data, mimeType } = node.inlineData; const processed = processInputAttachments({ type: "file", diff --git a/js/src/instrumentation/providers/groq-instrumentation.ts b/js/src/instrumentation/providers/groq-instrumentation.ts index c04f8d24d..610076cfe 100644 --- a/js/src/instrumentation/providers/groq-instrumentation.ts +++ b/js/src/instrumentation/providers/groq-instrumentation.ts @@ -7,7 +7,12 @@ import { isObject, SpanTypeAttribute, } from "../../../util/index"; -import { Attachment, withCurrent, type Span } from "../../logger"; +import { + _internalCaptureAttachmentsEnabled, + Attachment, + withCurrent, + type Span, +} from "../../logger"; import { convertDataToBlob, getExtensionFromMediaType, @@ -126,7 +131,10 @@ function extractGroqAudioInput( span: Span, ): { input: unknown; metadata: Record } { const source = params.file ?? params.url; - const filePart = groqAudioFilePart(source); + const filePart = groqAudioFilePart( + source, + _internalCaptureAttachmentsEnabled(span), + ); const input = { operation, ...(params.prompt !== undefined ? { prompt: params.prompt } : {}), @@ -149,11 +157,15 @@ function extractGroqAudioInput( }; } -function groqAudioFilePart(source: unknown): unknown | undefined { +function groqAudioFilePart( + source: unknown, + captureAttachments: boolean, +): unknown | undefined { if (source === undefined || source === null) return undefined; if (typeof source === "string") { if (source.startsWith("data:")) { + if (!captureAttachments) return undefined; const contentType = source.match(/^data:([^;,]+)/)?.[1]; const blob = contentType ? convertDataToBlob(source, contentType) : null; if (contentType && blob) { @@ -175,7 +187,7 @@ function groqAudioFilePart(source: unknown): unknown | undefined { }; } - if (!isObject(source)) return undefined; + if (!captureAttachments || !isObject(source)) return undefined; const filename = typeof source.name === "string" && source.name ? filenameFromPath(source.name) @@ -209,7 +221,7 @@ function observeGroqAudioInput( input: Record, span: Span, ): void { - if (!isObject(source)) return; + if (!_internalCaptureAttachmentsEnabled(span) || !isObject(source)) return; const body = isObject(source.body) ? source.body @@ -361,6 +373,7 @@ function captureGroqSpeechResponse( startTime: number, ): boolean { if (!isObject(response)) return false; + const captureAttachments = _internalCaptureAttachmentsEnabled(span); const headerContentType = responseHeader(response, "content-type")?.split( ";", @@ -399,7 +412,7 @@ function captureGroqSpeechResponse( }); firstChunk = false; } - chunks.push(new Uint8Array(chunk)); + if (captureAttachments) chunks.push(new Uint8Array(chunk)); }; const complete = () => { if (finished) return; diff --git a/js/src/instrumentation/providers/ollama-instrumentation.test.ts b/js/src/instrumentation/providers/ollama-instrumentation.test.ts index 4503d4ddf..74bad0b46 100644 --- a/js/src/instrumentation/providers/ollama-instrumentation.test.ts +++ b/js/src/instrumentation/providers/ollama-instrumentation.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, it, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import iso from "../../isomorph"; import { configureNode } from "../../node/config"; import { @@ -12,6 +12,8 @@ import { } from "./ollama-instrumentation"; configureNode(); +beforeEach(() => vi.stubEnv("BRAINTRUST_CAPTURE_ATTACHMENTS", "true")); +afterEach(() => vi.unstubAllEnvs()); describe("Ollama instrumentation extraction", () => { it("normalizes chat inputs, tools, and supported request metadata", () => { diff --git a/js/src/instrumentation/providers/ollama-instrumentation.ts b/js/src/instrumentation/providers/ollama-instrumentation.ts index 405b3dff9..17d5c23c5 100644 --- a/js/src/instrumentation/providers/ollama-instrumentation.ts +++ b/js/src/instrumentation/providers/ollama-instrumentation.ts @@ -1,6 +1,6 @@ import { SpanTypeAttribute, isObject } from "../../../util/index"; import iso from "../../isomorph"; -import { Attachment } from "../../logger"; +import { _internalCaptureAttachmentsEnabled, Attachment } from "../../logger"; import { processInputAttachments } from "../../wrappers/attachment-utils"; import type { OllamaChatRequest, @@ -199,7 +199,17 @@ function normalizeTextAndImages( const imageParts: Record[] = []; const unrecognizedImages: unknown[] = []; + const captureAttachments = _internalCaptureAttachmentsEnabled(); for (const image of images) { + if (!captureAttachments) { + if ( + image instanceof URL || + (typeof image === "string" && /^https?:\/\//i.test(image)) + ) { + imageParts.push({ type: "image_url", image_url: { url: image } }); + } + continue; + } let localImagePath: string | undefined; let localPathMediaType: string | undefined; if ( diff --git a/js/src/instrumentation/providers/openai-agents-trace-processor.ts b/js/src/instrumentation/providers/openai-agents-trace-processor.ts index a335f2efb..42d5c32b4 100644 --- a/js/src/instrumentation/providers/openai-agents-trace-processor.ts +++ b/js/src/instrumentation/providers/openai-agents-trace-processor.ts @@ -6,7 +6,10 @@ import { NOOP_SPAN, currentSpan, startSpan as startBaseSpan, + withCurrent, } from "../../logger"; +import { processInputAttachments } from "../../wrappers/attachment-utils"; +import { processImagesInOutput } from "./openai-span-data"; import { INSTRUMENTATION_NAMES, withSpanInstrumentationName, @@ -325,7 +328,9 @@ export class OpenAIAgentsTraceProcessor { return Promise.resolve(); } - const logData = this.extractLogData(span); + const logData = withCurrent(braintrustSpan, () => + this.extractLogData(span), + ); braintrustSpan.log({ error: span.error, ...logData, @@ -415,8 +420,8 @@ export class OpenAIAgentsTraceProcessor { }; return { - input: spanData._input, - output, + input: processInputAttachments(spanData._input), + output: processInputAttachments(processImagesInOutput(output)), metadata: isObject(response) ? this.omitKeys(response, ["output", "usage"]) : {}, @@ -459,8 +464,8 @@ export class OpenAIAgentsTraceProcessor { span: OpenAIAgentsSpan, ): Record { return { - input: spanData.input, - output: spanData.output, + input: processInputAttachments(spanData.input), + output: processInputAttachments(processImagesInOutput(spanData.output)), metadata: { model: spanData.model, model_config: spanData.model_config, @@ -493,7 +498,10 @@ export class OpenAIAgentsTraceProcessor { spanData: OpenAIAgentsTranscriptionSpanData, ): Record { return { - input: spanData.input, + input: processInputAttachments({ + type: "input_audio", + input_audio: spanData.input, + })?.input_audio, output: spanData.output, metadata: { model: spanData.model, @@ -507,7 +515,10 @@ export class OpenAIAgentsTraceProcessor { ): Record { return { input: spanData.input, - output: spanData.output, + output: processInputAttachments({ + type: "input_audio", + input_audio: spanData.output, + })?.input_audio, metadata: { model: spanData.model, model_config: spanData.model_config, diff --git a/js/src/instrumentation/providers/openai-instrumentation.test.ts b/js/src/instrumentation/providers/openai-instrumentation.test.ts index 90e162c6b..6827b88fe 100644 --- a/js/src/instrumentation/providers/openai-instrumentation.test.ts +++ b/js/src/instrumentation/providers/openai-instrumentation.test.ts @@ -1316,9 +1316,21 @@ describe("processImagesInOutput", () => { }; expect(processImagesInOutput(output, false)).toEqual({ - ...output, - result: "", + type: "image_generation_call", + output_format: "png", + revised_prompt: "A red pixel", }); + expect( + processImagesInOutput( + [ + { type: "image_generation_call", result: "AQID" }, + { type: "message", content: [{ type: "output_text", text: "done" }] }, + ], + false, + ), + ).toEqual([ + { type: "message", content: [{ type: "output_text", text: "done" }] }, + ]); }); describe("image_generation_call conversion", () => { @@ -1334,7 +1346,7 @@ describe("processImagesInOutput", () => { revised_prompt: "A red pixel", }; - const result = processImagesInOutput(output); + const result = processImagesInOutput(output, true); expect(result.type).toBe("image_generation_call"); expect(result.output_format).toBe("png"); @@ -1356,7 +1368,7 @@ describe("processImagesInOutput", () => { output_format: "jpg", }; - const result = processImagesInOutput(output); + const result = processImagesInOutput(output, true); const attachment = result.result as Attachment; expect(attachment.reference.filename).toBe("generated_image.jpg"); @@ -1373,7 +1385,7 @@ describe("processImagesInOutput", () => { output_format: "webp", }; - const result = processImagesInOutput(output); + const result = processImagesInOutput(output, true); const attachment = result.result as Attachment; expect(attachment.reference.filename).toBe("generated_image.webp"); @@ -1389,7 +1401,7 @@ describe("processImagesInOutput", () => { result: base64Image, }; - const result = processImagesInOutput(output); + const result = processImagesInOutput(output, true); const attachment = result.result as Attachment; expect(attachment.reference.filename).toBe("generated_image.png"); @@ -1406,7 +1418,7 @@ describe("processImagesInOutput", () => { revised_prompt: "A beautiful sunset", }; - const result = processImagesInOutput(output); + const result = processImagesInOutput(output, true); const attachment = result.result as Attachment; expect(attachment.reference.filename).toContain("A_beautiful_sunset"); @@ -1425,7 +1437,7 @@ describe("processImagesInOutput", () => { revised_prompt: longPrompt, }; - const result = processImagesInOutput(output); + const result = processImagesInOutput(output, true); const attachment = result.result as Attachment; // Should be truncated to 50 chars + sanitized + extension @@ -1442,7 +1454,7 @@ describe("processImagesInOutput", () => { revised_prompt: "Hello@World#Test!", }; - const result = processImagesInOutput(output); + const result = processImagesInOutput(output, true); const attachment = result.result as Attachment; expect(attachment.reference.filename).toBe("Hello_World_Test_.png"); @@ -1458,7 +1470,7 @@ describe("processImagesInOutput", () => { output_format: "jpg", }; - const result = processImagesInOutput(output); + const result = processImagesInOutput(output, true); const attachment = result.result as Attachment; expect(attachment.reference.filename).toBe("generated_image.jpg"); @@ -1477,7 +1489,7 @@ describe("processImagesInOutput", () => { another_field: 123, }; - const result = processImagesInOutput(output); + const result = processImagesInOutput(output, true); expect(result.type).toBe("image_generation_call"); expect(result.output_format).toBe("png"); @@ -1501,7 +1513,7 @@ describe("processImagesInOutput", () => { { type: "text", content: "World" }, ]; - const result = processImagesInOutput(output); + const result = processImagesInOutput(output, true); expect(Array.isArray(result)).toBe(true); expect(result[0]).toEqual({ type: "text", content: "Hello" }); @@ -1522,7 +1534,7 @@ describe("processImagesInOutput", () => { ], ]; - const result = processImagesInOutput(output); + const result = processImagesInOutput(output, true); expect(Array.isArray(result)).toBe(true); expect(Array.isArray(result[0])).toBe(true); @@ -1546,7 +1558,7 @@ describe("processImagesInOutput", () => { }, ]; - const result = processImagesInOutput(output); + const result = processImagesInOutput(output, true); expect(result[0].result).toBeInstanceOf(Attachment); expect(result[1].result).toBeInstanceOf(Attachment); @@ -1568,7 +1580,7 @@ describe("processImagesInOutput", () => { content: "Hello world", }; - const result = processImagesInOutput(output); + const result = processImagesInOutput(output, true); expect(result).toEqual(output); }); @@ -1578,7 +1590,7 @@ describe("processImagesInOutput", () => { data: "some data", }; - const result = processImagesInOutput(output); + const result = processImagesInOutput(output, true); expect(result).toEqual(output); }); @@ -1588,7 +1600,7 @@ describe("processImagesInOutput", () => { content: "Generated text", }; - const result = processImagesInOutput(output); + const result = processImagesInOutput(output, true); expect(result).toEqual(output); }); @@ -1621,7 +1633,7 @@ describe("processImagesInOutput", () => { result: null, }; - const result = processImagesInOutput(output); + const result = processImagesInOutput(output, true); expect(result).toEqual(output); }); @@ -1631,7 +1643,7 @@ describe("processImagesInOutput", () => { result: undefined, }; - const result = processImagesInOutput(output); + const result = processImagesInOutput(output, true); expect(result).toEqual(output); }); }); @@ -1643,7 +1655,7 @@ describe("processImagesInOutput", () => { result: 12345, // Not a string }; - const result = processImagesInOutput(output); + const result = processImagesInOutput(output, true); expect(result).toEqual(output); expect(result.result).toBe(12345); }); @@ -1654,7 +1666,7 @@ describe("processImagesInOutput", () => { result: { data: "some data" }, }; - const result = processImagesInOutput(output); + const result = processImagesInOutput(output, true); expect(result).toEqual(output); }); @@ -1664,7 +1676,7 @@ describe("processImagesInOutput", () => { result: ["data"], }; - const result = processImagesInOutput(output); + const result = processImagesInOutput(output, true); expect(result).toEqual(output); }); @@ -1674,7 +1686,7 @@ describe("processImagesInOutput", () => { result: true, }; - const result = processImagesInOutput(output); + const result = processImagesInOutput(output, true); expect(result).toEqual(output); }); @@ -1688,7 +1700,7 @@ describe("processImagesInOutput", () => { revised_prompt: 12345, // Not a string }; - const result = processImagesInOutput(output); + const result = processImagesInOutput(output, true); const attachment = result.result as Attachment; expect(attachment.reference.filename).toBe("generated_image.png"); diff --git a/js/src/instrumentation/providers/openai-instrumentation.ts b/js/src/instrumentation/providers/openai-instrumentation.ts index 521c5e553..49b3e54bf 100644 --- a/js/src/instrumentation/providers/openai-instrumentation.ts +++ b/js/src/instrumentation/providers/openai-instrumentation.ts @@ -1,3 +1,5 @@ +import { _internalCaptureAttachmentsEnabled } from "../../logger"; +import { processInputAttachments } from "../../wrappers/attachment-utils"; import { interceptOpenAIMedia } from "./openai-media"; import { traceAsyncChannel, @@ -71,9 +73,11 @@ export function registerOpenAIInstrumentation(): void { name: "Chat Completion", type: SpanTypeAttribute.LLM, extractInput: ([params]) => extractOpenAIChatInput(params), - extractOutput: (result) => { - return result?.choices; - }, + extractOutput: (result, event) => + processChatCompletionOutput( + result?.choices, + event?.arguments?.[0].audio?.format, + ), extractMetrics: (result, startTime, endEvent) => { const metrics = withCachedMetric( parseMetricsFromUsage(result?.usage), @@ -119,9 +123,11 @@ export function registerOpenAIInstrumentation(): void { name: "Chat Completion", type: SpanTypeAttribute.LLM, extractInput: ([params]) => extractOpenAIChatInput(params), - extractOutput: (result) => { - return result?.choices; - }, + extractOutput: (result, event) => + processChatCompletionOutput( + result?.choices, + event?.arguments?.[0].audio?.format, + ), extractMetrics: (result, startTime, endEvent) => { const metrics = withCachedMetric( parseMetricsFromUsage(result?.usage), @@ -246,6 +252,37 @@ export function registerOpenAIInstrumentation(): void { }); } +function processChatCompletionOutput( + choices: OpenAIChatChoice[] | undefined, + audioFormat: string | undefined, +): unknown { + const captureAttachments = _internalCaptureAttachmentsEnabled(); + return choices?.map((choice) => ({ + ...choice, + // Parsed outputs and tool arguments are application JSON, not media parts. + message: { + ...choice.message, + ...(Array.isArray(choice.message.content) + ? { + content: processInputAttachments( + choice.message.content, + captureAttachments, + ), + } + : {}), + ...(choice.message.audio + ? { + audio: processInputAttachments( + { role: "assistant", audio: choice.message.audio }, + captureAttachments, + audioFormat, + ).audio, + } + : {}), + }, + })); +} + function getCachedMetricFromEndEvent(endEvent: unknown): number | undefined { if (!isObject(endEvent)) { return undefined; diff --git a/js/src/instrumentation/providers/openai-media.ts b/js/src/instrumentation/providers/openai-media.ts index 3e20aec1f..3bb8dda73 100644 --- a/js/src/instrumentation/providers/openai-media.ts +++ b/js/src/instrumentation/providers/openai-media.ts @@ -1,4 +1,10 @@ -import { Attachment, startSpan, withCurrent, type Span } from "../../logger"; +import { + _internalCaptureAttachmentsEnabled, + Attachment, + startSpan, + withCurrent, + type Span, +} from "../../logger"; import { debugLogger } from "../../debug-logger"; import { getCurrentUnixTimestamp } from "../../util"; import { @@ -8,7 +14,6 @@ import { import { convertDataToBlob, getExtensionFromMediaType, - isAutoCaptureAttachmentsEnabled, } from "../../wrappers/attachment-utils"; import { isObject } from "../../../util/index"; import { isAsyncIterable, patchStreamIfNeeded } from "../core/stream-patcher"; @@ -661,7 +666,7 @@ export function interceptOpenAIMedia( return channel.intercept((target, thisArg, args, additional) => { if (isAutoInstrumentationSuppressed()) return target.apply(thisArg, args); const params = args[0]; - const captureAttachments = isAutoCaptureAttachmentsEnabled(); + const captureAttachments = _internalCaptureAttachmentsEnabled(); let span: Span; try { const { name, spanAttributes, spanInfoMetadata } = buildStartSpanArgs( diff --git a/js/src/instrumentation/providers/openai-span-data.ts b/js/src/instrumentation/providers/openai-span-data.ts index a1414ca84..9e5635dc8 100644 --- a/js/src/instrumentation/providers/openai-span-data.ts +++ b/js/src/instrumentation/providers/openai-span-data.ts @@ -1,7 +1,7 @@ /* eslint-disable @typescript-eslint/no-explicit-any */ -import { Attachment } from "../../logger"; +import { _internalCaptureAttachmentsEnabled, Attachment } from "../../logger"; import { - isAutoCaptureAttachmentsEnabled, + omitMediaData, processInputAttachments, } from "../../wrappers/attachment-utils"; import { isObject } from "../../../util/index"; @@ -58,7 +58,7 @@ export function extractOpenAIBatchInput( const input = endpoint === "/v1/chat/completions" ? params.messages : params.input; return { - input: processInputAttachments(input, isAutoCaptureAttachmentsEnabled()), + input: processInputAttachments(input, _internalCaptureAttachmentsEnabled()), metadata: batchMetadata(params), }; } @@ -69,7 +69,10 @@ export function extractOpenAIChatInput(params: Record): { } { const { messages, ...metadata } = params; return { - input: processInputAttachments(messages, isAutoCaptureAttachmentsEnabled()), + input: processInputAttachments( + messages, + _internalCaptureAttachmentsEnabled(), + ), metadata: { ...metadata, provider: "openai" }, }; } @@ -80,7 +83,7 @@ export function extractOpenAIResponsesInput(params: Record): { } { const { input, ...metadata } = params; return { - input: processInputAttachments(input, isAutoCaptureAttachmentsEnabled()), + input: processInputAttachments(input, _internalCaptureAttachmentsEnabled()), metadata: { ...metadata, provider: "openai" }, }; } @@ -98,12 +101,12 @@ export function extractOpenAIResponsesMetadata( /** Convert Responses API base64 image outputs to Braintrust attachments. */ export function processImagesInOutput( output: any, - captureAttachments = isAutoCaptureAttachmentsEnabled(), + captureAttachments = _internalCaptureAttachmentsEnabled(), ): any { if (Array.isArray(output)) { - return output.map((item) => - processImagesInOutput(item, captureAttachments), - ); + return output + .map((item) => processImagesInOutput(item, captureAttachments)) + .filter((item) => item !== undefined); } if ( @@ -113,7 +116,7 @@ export function processImagesInOutput( output.result ) { if (!captureAttachments) { - return { ...output, result: "" }; + return omitMediaData(output, "result"); } const fileExtension = output.output_format || "png"; const contentType = `image/${fileExtension}`; diff --git a/js/src/instrumentation/providers/pi-coding-agent-instrumentation.test.ts b/js/src/instrumentation/providers/pi-coding-agent-instrumentation.test.ts index 43d9cedb0..2002880c3 100644 --- a/js/src/instrumentation/providers/pi-coding-agent-instrumentation.test.ts +++ b/js/src/instrumentation/providers/pi-coding-agent-instrumentation.test.ts @@ -25,11 +25,18 @@ vi.mock("../../isomorph", async (importOriginal) => { }; }); -vi.mock("../../logger", () => ({ - withCurrent: (_span: unknown, callback: () => unknown) => callback(), - startSpan: (...args: unknown[]) => mockStartSpan(...args), - _internalStartSpan: (...args: unknown[]) => mockStartSpan(...args), -})); +vi.mock("../../logger", async () => { + const { default: iso } = await import("../../isomorph"); + return { + _internalCaptureAttachmentsEnabled: () => + iso.getEnv("BRAINTRUST_CAPTURE_ATTACHMENTS") === "true", + CAPTURE_ATTACHMENTS: Symbol.for("braintrust.captureAttachments"), + BaseAttachment: class {}, + _internalStartSpan: (...args: unknown[]) => mockStartSpan(...args), + startSpan: (...args: unknown[]) => mockStartSpan(...args), + withCurrent: (_span: unknown, callback: () => unknown) => callback(), + }; +}); import { isAutoInstrumentationSuppressed } from "../auto-instrumentation-suppression"; import { registerPiCodingAgentInstrumentation } from "./pi-coding-agent-instrumentation"; diff --git a/js/src/instrumentation/providers/pi-coding-agent-instrumentation.ts b/js/src/instrumentation/providers/pi-coding-agent-instrumentation.ts index bb1f83b41..481c1f841 100644 --- a/js/src/instrumentation/providers/pi-coding-agent-instrumentation.ts +++ b/js/src/instrumentation/providers/pi-coding-agent-instrumentation.ts @@ -2,7 +2,12 @@ import { toLoggedError } from "../core"; import type { ChannelMessage } from "../core/channel-definitions"; import iso, { type IsoAsyncLocalStorage } from "../../isomorph"; import { debugLogger } from "../../debug-logger"; -import { _internalStartSpan as startBaseSpan, withCurrent } from "../../logger"; +import { + _internalCaptureAttachmentsEnabled, + CAPTURE_ATTACHMENTS, + withCurrent, + _internalStartSpan as startBaseSpan, +} from "../../logger"; import type { Span } from "../../logger"; import { INSTRUMENTATION_NAMES, @@ -363,14 +368,19 @@ async function startPiLlmSpan( ...extractToolMetadata(context.tools), "pi_coding_agent.operation": `agent.${property}`, }; + const captureAttachments = _internalCaptureAttachmentsEnabled(state.span); const span = startBaseSpan( withSpanInstrumentationName( { event: { - input: processInputAttachments(normalizePiContextInput(context)), + input: processInputAttachments( + normalizePiContextInput(context), + captureAttachments, + ), metadata, }, name: getLlmSpanName(model), + [CAPTURE_ATTACHMENTS]: captureAttachments, parent: await state.span.export(), spanAttributes: { type: SpanTypeAttribute.LLM }, }, @@ -535,13 +545,19 @@ async function handlePiAgentEvent( switch (event.type) { case "message_end": if (isPiAssistantMessage(event.message)) { - state.output = extractAssistantOutput(event.message); + state.output = extractAssistantOutput( + event.message, + _internalCaptureAttachmentsEnabled(state.span), + ); } return; case "turn_end": state.turnEnded = true; if (isPiAssistantMessage(event.message)) { - state.output = extractAssistantOutput(event.message); + state.output = extractAssistantOutput( + event.message, + _internalCaptureAttachmentsEnabled(state.span), + ); if (!state.collectedLlmUsageMetrics) { addMetrics(state.metrics, extractUsageMetrics(event.message.usage)); } @@ -582,6 +598,7 @@ async function startPiToolSpan( metadata, }, name: event.toolName || "tool", + [CAPTURE_ATTACHMENTS]: _internalCaptureAttachmentsEnabled(state.span), parent: await state.span.export(), spanAttributes: { type: SpanTypeAttribute.TOOL }, }, @@ -689,7 +706,14 @@ function finishPiLlmSpan( ...(message ? extractAssistantMetadata(message) : {}), }, metrics, - ...(message ? { output: extractAssistantOutput(message) } : {}), + ...(message + ? { + output: extractAssistantOutput( + message, + _internalCaptureAttachmentsEnabled(llmState.span), + ), + } + : {}), }); } finally { llmState.span.end(); @@ -797,14 +821,20 @@ function normalizeToolCall(toolCall: PiToolCall): unknown { }; } -function extractAssistantOutput(message: PiAssistantMessage): unknown { - return processInputAttachments([ - { - finish_reason: normalizeStopReason(message.stopReason), - index: 0, - message: normalizeAssistantMessage(message), - }, - ]); +function extractAssistantOutput( + message: PiAssistantMessage, + captureAttachments: boolean, +): unknown { + return processInputAttachments( + [ + { + finish_reason: normalizeStopReason(message.stopReason), + index: 0, + message: normalizeAssistantMessage(message), + }, + ], + captureAttachments, + ); } function isPiUserMessage( diff --git a/js/src/instrumentation/providers/strands-agent-sdk-instrumentation.test.ts b/js/src/instrumentation/providers/strands-agent-sdk-instrumentation.test.ts index 8439682f6..b29a598ce 100644 --- a/js/src/instrumentation/providers/strands-agent-sdk-instrumentation.test.ts +++ b/js/src/instrumentation/providers/strands-agent-sdk-instrumentation.test.ts @@ -24,7 +24,9 @@ const { mockWithCurrent, mockNewAsyncLocalStorage, mockStartSpan } = vi.hoisted( vi.mock("../../isomorph", () => ({ default: { - getEnv: vi.fn(), + getEnv: vi.fn((name: string) => + name === "BRAINTRUST_CAPTURE_ATTACHMENTS" ? "true" : undefined, + ), newAsyncLocalStorage: mockNewAsyncLocalStorage, newTracingChannel: vi.fn(), }, @@ -39,6 +41,9 @@ vi.mock("../../logger", async (importOriginal) => { }; }); +vi.mock("../../lru-cache", { spy: true }); + +import { LRUCache } from "../../lru-cache"; import iso from "../../isomorph"; import { Attachment } from "../../logger"; import { isAutoInstrumentationSuppressed } from "../auto-instrumentation-suppression"; @@ -318,11 +323,18 @@ describe("registerStrandsAgentSDKInstrumentation", () => { }); it.each([ - ["binary objects", new Uint8Array([1, 2, 3])], - ["base64 strings", "AQID"], + ["binary objects", new Uint8Array([1, 2, 3]), true], + ["omitted binary objects", new Uint8Array([1, 2, 3]), false], + ["base64 strings", "AQID", true], + ["omitted base64 strings", "AQID", false], ])( "converts media from %s to one attachment shared by agent and model spans", - async (_description, bytes) => { + async (_description, bytes, captureAttachments) => { + vi.mocked(iso.getEnv).mockImplementation((name) => + name === "BRAINTRUST_CAPTURE_ATTACHMENTS" + ? String(captureAttachments) + : undefined, + ); registerStrandsAgentSDKInstrumentation(); const handlers = handlersByName.get( @@ -368,10 +380,16 @@ describe("registerStrandsAgentSDKInstrumentation", () => { (span) => span.args.name === "Strands model: gpt-4o-mini", ); const rootAttachment = - rootSpan?.args.event.input[0].document.source.bytes; + rootSpan?.args.event.input[0].document.source?.bytes; const modelAttachment = - modelSpan?.args.event.input[0].content[0].document.source.bytes; + modelSpan?.args.event.input[0].content[0].document.source?.bytes; + if (!captureAttachments) { + expect(LRUCache).not.toHaveBeenCalled(); + expect(rootAttachment).toBeUndefined(); + expect(modelAttachment).toBeUndefined(); + return; + } expect(rootAttachment).toBeInstanceOf(Attachment); expect(rootAttachment.reference).toMatchObject({ content_type: "application/pdf", diff --git a/js/src/instrumentation/providers/strands-agent-sdk-instrumentation.ts b/js/src/instrumentation/providers/strands-agent-sdk-instrumentation.ts index 03ea535e3..c3b791078 100644 --- a/js/src/instrumentation/providers/strands-agent-sdk-instrumentation.ts +++ b/js/src/instrumentation/providers/strands-agent-sdk-instrumentation.ts @@ -2,6 +2,7 @@ import { toLoggedError } from "../core"; import { isAsyncIterable, patchStreamIfNeeded } from "../core/stream-patcher"; import { debugLogger } from "../../debug-logger"; import { + _internalCaptureAttachmentsEnabled, Attachment, BaseAttachment, startSpan as startBaseSpan, @@ -15,7 +16,10 @@ import { import { LRUCache } from "../../lru-cache"; import { getCurrentUnixTimestamp } from "../../util"; import { SpanTypeAttribute, isObject } from "../../../util/index"; -import { convertDataToBlob } from "../../wrappers/attachment-utils"; +import { + convertDataToBlob, + omitMediaData, +} from "../../wrappers/attachment-utils"; import { runWithAutoInstrumentationSuppressed } from "../auto-instrumentation-suppression"; import { strandsAgentSDKChannels } from "./strands-agent-sdk-channels"; import type { @@ -91,10 +95,13 @@ type ActiveChildParents = WeakMap>; const MAX_STRANDS_STRING_ATTACHMENT_CACHE_ENTRIES = 32; -type StrandsAttachmentCache = { - objects: WeakMap>; - strings: LRUCache>; -}; +type StrandsAttachmentCache = + | { captureAttachments: false } + | { + captureAttachments: true; + objects: WeakMap>; + strings: LRUCache>; + }; class StrandsAgentSDKInstrumentationConsumer { private readonly activeChildParents: ActiveChildParents = new WeakMap(); @@ -235,7 +242,7 @@ function startAgentStream( const parentSpan = agent ? getOnlyChildParent(activeChildParents, agent) : undefined; - const attachmentCache = createStrandsAttachmentCache(); + const attachmentCache = createStrandsAttachmentCache(parentSpan); const processedInput = processStrandsInputAttachments(input, attachmentCache); const span = parentSpan ? withCurrent(parentSpan, () => @@ -777,7 +784,8 @@ function finalizeAgentStream( ...(output !== undefined ? { output } : {}), }); state.span.end(); - state.attachmentCache.strings.clear(); + if (state.attachmentCache.captureAttachments) + state.attachmentCache.strings.clear(); } function finalizeMultiAgentStream( @@ -1002,8 +1010,11 @@ const STRANDS_MEDIA_TYPES: Record = { xml: "application/xml", }; -function createStrandsAttachmentCache(): StrandsAttachmentCache { +function createStrandsAttachmentCache(parent?: Span): StrandsAttachmentCache { + if (!_internalCaptureAttachmentsEnabled(parent)) + return { captureAttachments: false }; return { + captureAttachments: true, objects: new WeakMap(), strings: new LRUCache({ max: MAX_STRANDS_STRING_ATTACHMENT_CACHE_ENTRIES, @@ -1130,35 +1141,44 @@ function createStrandsMediaAttachment( return undefined; } - const contentType = STRANDS_MEDIA_TYPES[format.toLowerCase()]; - if (!contentType) { - return undefined; - } - const filename = - mediaKey === "document" && - typeof media.name === "string" && - media.name.length > 0 - ? media.name - : `${mediaKey}.${format.toLowerCase()}`; - const attachment = getOrCreateStrandsAttachment( - source.bytes, - filename, - contentType, - cache, - ); - if (!attachment) { - return undefined; + let attachment: Attachment | undefined; + if (cache.captureAttachments) { + const contentType = STRANDS_MEDIA_TYPES[format.toLowerCase()]; + if (!contentType) { + return undefined; + } + const filename = + mediaKey === "document" && + typeof media.name === "string" && + media.name.length > 0 + ? media.name + : `${mediaKey}.${format.toLowerCase()}`; + attachment = getOrCreateStrandsAttachment( + source.bytes, + filename, + contentType, + cache, + ); + if (!attachment) { + return undefined; + } } - const { type: _type, ...serializedMedia } = media; - const { type: _sourceType, ...serializedSource } = source; + const { type: _type, source: _source, ...serializedMedia } = media; + const serializedSource = cache.captureAttachments + ? omitMediaData(source, "type") + : omitMediaData(source, "type", "bytes"); return { [mediaKey]: { ...serializedMedia, - source: { - ...serializedSource, - bytes: attachment, - }, + ...(serializedSource || attachment + ? { + source: { + ...serializedSource, + ...(attachment ? { bytes: attachment } : {}), + }, + } + : {}), }, }; } @@ -1167,7 +1187,7 @@ function getOrCreateStrandsAttachment( data: unknown, filename: string, contentType: string, - cache: StrandsAttachmentCache, + cache: Extract, ): Attachment | undefined { const key = `${contentType}\0${filename}`; const attachments = diff --git a/js/src/instrumentation/providers/voyageai-instrumentation.test.ts b/js/src/instrumentation/providers/voyageai-instrumentation.test.ts index 2cdb491b4..874b76472 100644 --- a/js/src/instrumentation/providers/voyageai-instrumentation.test.ts +++ b/js/src/instrumentation/providers/voyageai-instrumentation.test.ts @@ -21,6 +21,7 @@ describe("registerVoyageAIInstrumentation", () => { beforeEach(() => { backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); initLogger({ + captureAttachments: true, projectName: "voyageai-instrumentation.test.ts", projectId: "test-project-id", }); diff --git a/js/src/logger.ts b/js/src/logger.ts index 8ebaddaa4..a49594be1 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -753,6 +753,28 @@ function normalizeProxyConnUrl(proxyUrl: string): string { : proxyUrl; } +/** @internal */ +export const CAPTURE_ATTACHMENTS = Symbol.for("braintrust.captureAttachments"); + +function captureAttachmentsFromEnv(): boolean { + return ["1", "true"].includes( + iso.getEnv("BRAINTRUST_CAPTURE_ATTACHMENTS")?.trim().toLowerCase() ?? "", + ); +} + +/** + * @internal Resolve automatic attachment capture for a logger or span. Only plain + * fields are read, so parents and global state from other SDK bundles also work. + */ +export function _internalCaptureAttachmentsEnabled(parent?: object): boolean { + // Global state is unset until a platform entrypoint configures the SDK. + const owner = parent ?? (_globalState ? getSpanParentObject() : undefined); + const policy = owner && Reflect.get(owner, CAPTURE_ATTACHMENTS); + return typeof policy === "boolean" + ? policy + : (_globalState?.captureAttachments ?? captureAttachmentsFromEnv()); +} + export class BraintrustState { public id: string; public currentExperiment: Experiment | undefined; @@ -775,6 +797,8 @@ export class BraintrustState { public proxyUrl: string | null = null; public loggedIn: boolean = false; public gitMetadataSettings?: GitMetadataSettings; + /** @internal Default for instrumentation without a logger-specific policy. */ + public captureAttachments?: boolean; public debugLogLevel?: DebugLogLevel; private debugLogLevelConfigured = false; @@ -2886,6 +2910,7 @@ export class Logger implements Exportable { private lastStartTime: number; private lazyId: LazyValue; private calledStartSpan: boolean; + private readonly captureAttachments: boolean | undefined; // For type identification. public kind = "logger" as const; @@ -2896,6 +2921,7 @@ export class Logger implements Exportable { options: { computeMetadataArgs?: ProjectMetadataArgs; linkArgs?: LinkArgs; + captureAttachments?: boolean; } = {}, ) { this.lazyMetadata = lazyMetadata; @@ -2904,6 +2930,7 @@ export class Logger implements Exportable { this.lastStartTime = getCurrentUnixTimestamp(); this.lazyId = new LazyValue(() => this.id); this.calledStartSpan = false; + this.captureAttachments = options.captureAttachments; this.state = state; } @@ -2927,6 +2954,11 @@ export class Logger implements Exportable { return this.state; } + /** @internal */ + public get [CAPTURE_ATTACHMENTS](): boolean { + return this.captureAttachments ?? captureAttachmentsFromEnv(); + } + private parentObjectType() { return SpanObjectTypeV3.PROJECT_LOGS; } @@ -3004,6 +3036,7 @@ export class Logger implements Exportable { private startSpanImpl(args?: StartSpanArgs): Span { return new SpanImpl({ + [CAPTURE_ATTACHMENTS]: this[CAPTURE_ATTACHMENTS], ...args, // Sometimes `args` gets passed directly into this function, and it contains an undefined value for `state`. // To ensure that we always use this logger's state, we override the `state` argument no matter what. @@ -4976,6 +5009,15 @@ async function computeLoggerMetadata( } export type InitLoggerOptions = FullLoginOptions & { + /** + * Capture inline media as attachments in instrumentation. Explicit booleans override + * BRAINTRUST_CAPTURE_ATTACHMENTS (1 or true); capture is otherwise disabled. + * With setCurrent: false, this logger and its child spans have an independent policy + * and an omitted option uses the environment, not the global logger's setting. + * Current loggers also update the state's default when this option is supplied; + * omitting it preserves that default. Explicitly logged attachments are unaffected. + */ + captureAttachments?: boolean; projectName?: string; projectId?: string; _createInProjectGroup?: string; @@ -5048,11 +5090,18 @@ export function initLogger(options: Readonly = {}): Logger { }, ); + const setCurrent = options.setCurrent ?? true; + if (setCurrent && options.captureAttachments !== undefined) { + state.captureAttachments = options.captureAttachments; + } const ret = new Logger(state, lazyMetadata, { + captureAttachments: + options.captureAttachments ?? + (setCurrent ? state.captureAttachments : undefined), computeMetadataArgs, linkArgs, }); - if (options.setCurrent ?? true) { + if (setCurrent) { state.currentLogger = ret; } return ret; @@ -8006,6 +8055,8 @@ function _resolveSpanIds( * We suggest using one of the various `traced` methods, instead of creating Spans directly. See {@link Span.startSpan} for full details. */ export class SpanImpl implements Span { + /** @internal Snapshot inherited by descendants and delayed instrumentation. */ + public readonly [CAPTURE_ATTACHMENTS]: boolean; private _state: BraintrustState; private isMerge: boolean; @@ -8041,11 +8092,16 @@ export class SpanImpl implements Span { defaultRootType?: SpanType; spanId?: string; propagatedState?: PropagatedState | undefined; + [CAPTURE_ATTACHMENTS]?: boolean; } & Omit & InitialSpanWriteAsMergeArg & InternalSpanContextArg, ) { this._state = args.state; + this[CAPTURE_ATTACHMENTS] = + args[CAPTURE_ATTACHMENTS] ?? + args.state.captureAttachments ?? + captureAttachmentsFromEnv(); this._propagatedState = args.propagatedState; const instrumentationName = getSpanInstrumentationName(args) ?? @@ -8275,6 +8331,7 @@ export class SpanImpl implements Span { : { spanId: this._spanId, rootSpanId: this._rootSpanId }; return new SpanImpl({ state: this._state, + [CAPTURE_ATTACHMENTS]: this[CAPTURE_ATTACHMENTS], ...args, ...startSpanParentArgs({ state: this._state, @@ -8300,6 +8357,7 @@ export class SpanImpl implements Span { }; return new SpanImpl({ state: this._state, + [CAPTURE_ATTACHMENTS]: this[CAPTURE_ATTACHMENTS], ...args, ...startSpanParentArgs({ state: this._state, diff --git a/js/src/vendor-sdk-types/openai-common.ts b/js/src/vendor-sdk-types/openai-common.ts index a6ec770de..177f0f025 100644 --- a/js/src/vendor-sdk-types/openai-common.ts +++ b/js/src/vendor-sdk-types/openai-common.ts @@ -22,6 +22,7 @@ export interface OpenAIWithResponse { export interface OpenAIChatCreateParams { messages: unknown; stream?: boolean | null; + audio?: { format: string; [key: string]: unknown } | null; [key: string]: unknown; } diff --git a/js/src/wrappers/ai-sdk/normalize-logged-output.ts b/js/src/wrappers/ai-sdk/normalize-logged-output.ts index 40773f126..d60262fc5 100644 --- a/js/src/wrappers/ai-sdk/normalize-logged-output.ts +++ b/js/src/wrappers/ai-sdk/normalize-logged-output.ts @@ -1,4 +1,5 @@ import type { AISDKResult } from "../../vendor-sdk-types/ai-sdk"; +import { BaseAttachment } from "../../logger"; const REMOVE_NORMALIZED_VALUE = Symbol("braintrust.ai-sdk.remove-normalized"); @@ -21,7 +22,12 @@ function normalizeAISDKLoggedValue( .filter((entry) => entry !== REMOVE_NORMALIZED_VALUE); } - if (!value || typeof value !== "object") { + if ( + !value || + typeof value !== "object" || + value instanceof BaseAttachment || + value instanceof URL + ) { return value; } diff --git a/js/src/wrappers/ai-sdk/telemetry.ts b/js/src/wrappers/ai-sdk/telemetry.ts index 72c860c3d..ef4f328ce 100644 --- a/js/src/wrappers/ai-sdk/telemetry.ts +++ b/js/src/wrappers/ai-sdk/telemetry.ts @@ -365,7 +365,9 @@ export function braintrustAISDKTelemetry(): any { span.log({ ...(shouldRecordOutputs(event) ? { - output: processAISDKOutput(result, DEFAULT_DENY_OUTPUT_PATHS), + output: withCurrent(span, () => + processAISDKOutput(result, DEFAULT_DENY_OUTPUT_PATHS), + ), } : {}), metrics: extractTokenMetrics(result), @@ -464,7 +466,9 @@ export function braintrustAISDKTelemetry(): any { state.span.log({ ...(shouldRecordOutputs(event) ? { - output: finishOutput(result, state.operationName), + output: withCurrent(state.span, () => + finishOutput(result, state.operationName), + ), } : {}), metrics, @@ -547,9 +551,11 @@ export function braintrustAISDKTelemetry(): any { if (shouldRecordInputs(event)) { const callInput = workflowAgentCallInput ?? operationInput(event, operationName); - const { input, outputPromise } = workflowAgent - ? processAISDKWorkflowAgentCallInput(callInput) - : processAISDKCallInput(callInput); + const { input, outputPromise } = withCurrent(span, () => + workflowAgent + ? processAISDKWorkflowAgentCallInput(callInput) + : processAISDKCallInput(callInput), + ); logPayload.input = input; const state = operations.get(operationKey); if (state) { @@ -590,12 +596,15 @@ export function braintrustAISDKTelemetry(): any { const operationName = state?.operationName ?? "generateText"; const callInput = operationInput(event, operationName); const workflowAgent = operationName === "WorkflowAgent.stream"; - const processedInput = + const processInput = () => shouldRecordInputs(event) && workflowAgent ? processAISDKWorkflowAgentModelCallInput(callInput).input : shouldRecordInputs(event) ? processAISDKCallInput(callInput).input : undefined; + const processedInput = state + ? withCurrent(state.span, processInput) + : processInput(); if ( workflowAgent && state?.ownsSpan && @@ -738,7 +747,9 @@ export function braintrustAISDKTelemetry(): any { span.log({ ...(shouldRecordOutputs(event) ? { - output: processAISDKOutput(result, DEFAULT_DENY_OUTPUT_PATHS), + output: withCurrent(span, () => + processAISDKOutput(result, DEFAULT_DENY_OUTPUT_PATHS), + ), } : {}), metrics: extractTokenMetrics(result), diff --git a/js/src/wrappers/attachment-utils.test.ts b/js/src/wrappers/attachment-utils.test.ts index 3994ecb5c..1ba582e27 100644 --- a/js/src/wrappers/attachment-utils.test.ts +++ b/js/src/wrappers/attachment-utils.test.ts @@ -1,9 +1,7 @@ -import { afterEach, describe, expect, it } from "vitest"; +import { afterEach, describe, expect, it, vi } from "vitest"; import iso from "../isomorph"; -import { - isAutoCaptureAttachmentsEnabled, - processInputAttachments, -} from "./attachment-utils"; +import { _internalCaptureAttachmentsEnabled, Attachment } from "../logger"; +import { processInputAttachments } from "./attachment-utils"; const originalGetEnv = iso.getEnv; @@ -11,11 +9,11 @@ afterEach(() => { iso.getEnv = originalGetEnv; }); -describe("isAutoCaptureAttachmentsEnabled", () => { +describe("_internalCaptureAttachmentsEnabled", () => { it.each(["1", "true", " TRUE "])('accepts "%s"', (value) => { iso.getEnv = () => value; - expect(isAutoCaptureAttachmentsEnabled()).toBe(true); + expect(_internalCaptureAttachmentsEnabled()).toBe(true); }); it.each([undefined, "0", "false", "yes", "on", "unexpected"])( @@ -23,7 +21,7 @@ describe("isAutoCaptureAttachmentsEnabled", () => { (value) => { iso.getEnv = () => value; - expect(isAutoCaptureAttachmentsEnabled()).toBe(false); + expect(_internalCaptureAttachmentsEnabled()).toBe(false); }, ); }); @@ -50,7 +48,7 @@ it("omits inline attachment data while preserving remote references", () => { expect(processInputAttachments(input, false)).toEqual([ { type: "image_url", - image_url: { detail: "low", url: "" }, + image_url: { detail: "low" }, }, { type: "image_url", @@ -58,7 +56,85 @@ it("omits inline attachment data while preserving remote references", () => { }, { type: "file", - file: { filename: "document.pdf", file_data: "" }, + file: { filename: "document.pdf" }, }, ]); }); + +it("skips data URL parsing for disabled media across supported input formats", () => { + const data = "data:application/octet-stream;base64,AQID"; + const input = [ + { type: "image_url", image_url: { url: data } }, + { type: "image_base64", imageBase64: data }, + { type: "video_base64", video_base64: data }, + { type: "file", file: { file_data: data } }, + { type: "image", image: data }, + { type: "file", data, mediaType: "application/pdf" }, + ]; + const parse = vi.spyOn(String.prototype, "match"); + let output; + let parseCalls; + try { + output = processInputAttachments(input, false); + parseCalls = parse.mock.calls.length; + } finally { + parse.mockRestore(); + } + expect(parseCalls).toBe(0); + expect(output).toEqual([{ type: "file", mediaType: "application/pdf" }]); + expect(JSON.stringify(output)).not.toContain("AQID"); + expect(input[0].image_url?.url).toBe(data); +}); + +it.each([false, true])( + "applies capture=%s to data URL objects while preserving remote URLs", + (captureAttachments) => { + const data = new URL("data:image/png;base64,AQID"); + const remote = new URL("https://example.com/image.png"); + const input = [ + data, + { type: "image", image: data }, + { type: "file", data, mediaType: "image/png" }, + { type: "image_url", image_url: { url: data, detail: "low" } }, + { type: "file", file: { file_data: data, filename: "image.png" } }, + remote, + { type: "image", image: remote }, + { type: "file", data: remote, mediaType: "image/png" }, + ]; + const decode = vi.spyOn(globalThis, "atob"); + try { + const output = processInputAttachments(input, captureAttachments); + if (captureAttachments) { + expect(output.slice(0, 5)).toEqual([ + expect.any(Attachment), + { type: "image", image: expect.any(Attachment) }, + { + type: "file", + data: expect.any(Attachment), + mediaType: "image/png", + }, + { + type: "image_url", + image_url: { url: expect.any(Attachment), detail: "low" }, + }, + { + type: "file", + file: { file_data: expect.any(Attachment), filename: "image.png" }, + }, + ]); + } else { + expect(decode).not.toHaveBeenCalled(); + expect(output.slice(0, -3)).toEqual([ + { type: "file", mediaType: "image/png" }, + { type: "image_url", image_url: { detail: "low" } }, + { type: "file", file: { filename: "image.png" } }, + ]); + } + expect(output.slice(-3)).toEqual(input.slice(-3)); + expect(JSON.stringify(output)).not.toContain("AQID"); + expect(data.href).toBe("data:image/png;base64,AQID"); + } finally { + decode.mockRestore(); + } + }, +); diff --git a/js/src/wrappers/attachment-utils.ts b/js/src/wrappers/attachment-utils.ts index e18b526d5..eb6123845 100644 --- a/js/src/wrappers/attachment-utils.ts +++ b/js/src/wrappers/attachment-utils.ts @@ -1,13 +1,29 @@ -import { Attachment } from "../logger"; -import iso from "../isomorph"; +import { + Attachment, + BaseAttachment, + _internalCaptureAttachmentsEnabled, +} from "../logger"; -const CAPTURE_ATTACHMENTS_ENV_VAR = "BRAINTRUST_CAPTURE_ATTACHMENTS"; - -export function isAutoCaptureAttachmentsEnabled(): boolean { - const value = iso.getEnv(CAPTURE_ATTACHMENTS_ENV_VAR); - return ( - value !== undefined && ["1", "true"].includes(value.trim().toLowerCase()) - ); +/** Remove media fields without reading their values, dropping type-only blocks. */ +export function omitMediaData( + value: object, + ...fields: string[] +): Record | undefined { + const result: Record = {}; + let hasContent = false; + for (const key of Object.keys(value)) { + if (fields.includes(key)) continue; + const item: unknown = Reflect.get(value, key); + if (item === undefined) continue; + Object.defineProperty(result, key, { + value: item, + enumerable: true, + configurable: true, + writable: true, + }); + if (key !== "type") hasContent = true; + } + return hasContent ? result : undefined; } /** @@ -27,6 +43,8 @@ export function getExtensionFromMediaType(mediaType: string): string { "audio/wav": "wav", "audio/webm": "webm", "audio/ogg": "ogg", + "audio/aac": "aac", + "audio/pcm": "pcm", "video/mp4": "mp4", "video/webm": "webm", "application/pdf": "pdf", @@ -44,6 +62,10 @@ export function getExtensionFromMediaType(mediaType: string): string { */ export function convertDataToBlob(data: any, mediaType: string): Blob | null { try { + if (data instanceof URL) { + if (data.protocol !== "data:") return null; + data = data.href; + } if (typeof data === "string") { // Could be base64, data URL, or regular URL if (data.startsWith("data:")) { @@ -86,11 +108,12 @@ export function convertDataToBlob(data: any, mediaType: string): Blob | null { /** * Process input to extract and convert image/file content parts to Attachments - * Similar to processImagesInOutput in oai_responses.ts - replaces data in-place + * Only pass provider content containers, never arbitrary tool arguments or JSON results. */ export function processInputAttachments( input: any, - captureAttachments = true, + captureAttachments = _internalCaptureAttachmentsEnabled(), + audioOutputFormat?: string, ): any { if (!input) { return input; @@ -125,23 +148,151 @@ export function processInputAttachments( const processNode = (node: any): any => { if (Array.isArray(node)) { - return node.map(processNode); + return node.map(processNode).filter((item) => item !== undefined); } if (!node || typeof node !== "object") { return node; } + if (node instanceof BaseAttachment || node instanceof Date) return node; + // Tool arguments and JSON results are application data, not content parts. + if (node.type === "tool-call") return node; + if (node.type === "tool-result") { + return node.output?.type === "content" && Array.isArray(node.output.value) + ? { + ...node, + output: { ...node.output, value: processNode(node.output.value) }, + } + : node; + } + if (node instanceof URL) { + if (node.protocol !== "data:") return node; + if (!captureAttachments) return undefined; + const mediaType = inferMediaTypeFromDataUrl( + node.href, + "application/octet-stream", + ); + return ( + toAttachment( + node, + mediaType, + `file.${getExtensionFromMediaType(mediaType)}`, + ) ?? node + ); + } + + // Responses API and Agents SDK image/file content parts. + if (node.type === "input_image" || node.type === "input_file") { + const field = + node.type === "input_file" + ? "file_data" in node + ? "file_data" + : "file" + : "image_url" in node + ? "image_url" + : "image"; + const data = node[field]; + if (data instanceof BaseAttachment) return node; + if (typeof data === "string" || data instanceof URL) { + if ( + (data instanceof URL && data.protocol !== "data:") || + /^https?:/i.test(String(data)) + ) + return node; + if (!captureAttachments) return omitMediaData(node, field); + const mediaType = inferMediaTypeFromDataUrl( + String(data), + node.type === "input_image" + ? "image/png" + : "application/octet-stream", + ); + const filename = + node.filename || `file.${getExtensionFromMediaType(mediaType)}`; + const attachment = toAttachment(data, mediaType, filename); + if (attachment) return { ...node, [field]: attachment }; + } + } + + if (node.type === "audio" && typeof node.audio === "string") { + if (!captureAttachments) return omitMediaData(node, "audio"); + const format = node.format || "wav"; + const attachment = toAttachment( + node.audio, + format === "mp3" ? "audio/mpeg" : `audio/${format}`, + `audio.${format}`, + ); + if (attachment) return { ...node, audio: attachment }; + } + + const audioField = node.type === "input_audio" ? "input_audio" : "audio"; + if ( + (node.type === "input_audio" || node.role === "assistant") && + node[audioField] && + typeof node[audioField] === "object" + ) { + const audio = node[audioField]; + if (audio.data instanceof BaseAttachment) return node; + if (!captureAttachments) + return omitMediaData({ + ...node, + [audioField]: omitMediaData(audio, "data"), + }); + const format = + audio.format || + (node.role === "assistant" ? audioOutputFormat : undefined) || + "wav"; + const mediaType = + format === "mp3" + ? "audio/mpeg" + : format === "opus" + ? "audio/ogg" + : format === "pcm16" + ? "audio/pcm" + : `audio/${format}`; + const attachment = toAttachment( + audio.data, + mediaType, + `audio.${getExtensionFromMediaType(mediaType)}`, + ); + if (attachment) + return { ...node, [audioField]: { ...audio, data: attachment } }; + } + + // AI SDK generateText/streamText file parts wrap a GeneratedFile instance. + if ( + node.type === "file" && + node.file && + typeof node.file === "object" && + ("base64" in node.file || "uint8Array" in node.file) + ) { + if (!captureAttachments) return omitMediaData(node, "file"); + const mediaType = node.file.mediaType || "application/octet-stream"; + const attachment = toAttachment( + node.file.base64 ?? node.file.uint8Array, + mediaType, + `file.${getExtensionFromMediaType(mediaType)}`, + ); + if (attachment) return { ...node, file: attachment }; + } + // OpenAI chat image_url content format if ( node.type === "image_url" && node.image_url && typeof node.image_url === "object" && - typeof node.image_url.url === "string" && - node.image_url.url.startsWith("data:") + ((typeof node.image_url.url === "string" && + node.image_url.url.startsWith("data:")) || + (node.image_url.url instanceof URL && + node.image_url.url.protocol === "data:")) ) { + if (!captureAttachments) + return omitMediaData({ + ...node, + image_url: omitMediaData(node.image_url, "url"), + }); const mediaType = inferMediaTypeFromDataUrl( - node.image_url.url, + String(node.image_url.url), "image/png", ); const filename = `image.${getExtensionFromMediaType(mediaType)}`; @@ -179,6 +330,7 @@ export function processInputAttachments( typeof voyageBase64Value === "string" && voyageBase64Value.startsWith("data:") ) { + if (!captureAttachments) return omitMediaData(node, voyageBase64Key); const mediaType = inferMediaTypeFromDataUrl( voyageBase64Value, node.type === "video_base64" ? "video/mp4" : "image/png", @@ -201,11 +353,18 @@ export function processInputAttachments( node.type === "file" && node.file && typeof node.file === "object" && - typeof node.file.file_data === "string" && - node.file.file_data.startsWith("data:") + ((typeof node.file.file_data === "string" && + !/^https?:/i.test(node.file.file_data)) || + (node.file.file_data instanceof URL && + node.file.file_data.protocol === "data:")) ) { + if (!captureAttachments) + return omitMediaData({ + ...node, + file: omitMediaData(node.file, "file_data"), + }); const mediaType = inferMediaTypeFromDataUrl( - node.file.file_data, + String(node.file.file_data), "application/octet-stream", ); const filename = @@ -229,9 +388,20 @@ export function processInputAttachments( // AI SDK image content format if (node.type === "image" && node.image) { + if (node.image instanceof BaseAttachment) { + attachmentIndex++; + return node; + } + if (!captureAttachments) { + return (node.image instanceof URL && node.image.protocol !== "data:") || + (typeof node.image === "string" && /^https?:/i.test(node.image)) + ? node + : omitMediaData(node, "image"); + } let mediaType = "image/png"; - if (typeof node.image === "string" && node.image.startsWith("data:")) { - mediaType = inferMediaTypeFromDataUrl(node.image, mediaType); + const image = node.image instanceof URL ? node.image.href : node.image; + if (typeof image === "string" && image.startsWith("data:")) { + mediaType = inferMediaTypeFromDataUrl(image, mediaType); } else if (node.mediaType) { mediaType = node.mediaType; } @@ -253,8 +423,21 @@ export function processInputAttachments( } } - // AI SDK file content format - if (node.type === "file" && node.data) { + // AI SDK file parts and inline media in tool-result content. + if ( + ["file", "image-data", "file-data", "media"].includes(node.type) && + node.data + ) { + if (node.data instanceof BaseAttachment) { + attachmentIndex++; + return node; + } + if (!captureAttachments) { + return (node.data instanceof URL && node.data.protocol !== "data:") || + (typeof node.data === "string" && /^https?:/i.test(node.data)) + ? node + : omitMediaData(node, "data"); + } const mediaType = node.mediaType || "application/octet-stream"; const filename = node.filename || @@ -275,35 +458,16 @@ export function processInputAttachments( } } - // AI SDK tool output content format - if ( - (node.type === "image-data" || node.type === "file-data") && - node.data - ) { - const mediaType = node.mediaType || "application/octet-stream"; - const filename = - node.filename || - `output_${node.type === "image-data" ? "image" : "file"}_${attachmentIndex}.${getExtensionFromMediaType(mediaType)}`; - const attachment = toAttachment(node.data, mediaType, filename); - - if (attachment) { - attachmentIndex++; - return { - ...node, - data: attachment, - }; - } - } - const processed: Record = {}; for (const [key, value] of Object.entries(node)) { - processed[key] = processNode(value); + const result = processNode(value); + if (result !== undefined) processed[key] = result; } return processed; }; if (Array.isArray(input)) { - return input.map(processNode); + return input.map(processNode).filter((item) => item !== undefined); } return processNode(input); diff --git a/js/src/wrappers/groq.test.ts b/js/src/wrappers/groq.test.ts index d44daf491..5a3c19880 100644 --- a/js/src/wrappers/groq.test.ts +++ b/js/src/wrappers/groq.test.ts @@ -30,6 +30,7 @@ describe("groq wrapper", () => { beforeEach(() => { backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); initLogger({ + captureAttachments: true, projectId: "test-project-id", projectName: "groq.test.ts", }); @@ -426,39 +427,54 @@ describe("groq wrapper", () => { }); }); - test("captures speech consumed through a cloned response", async () => { - const wrapped = wrapGroq({ - audio: { - speech: { - create: vi.fn(async (_request: Record) => - Promise.resolve( - new Response(new Uint8Array([1, 2, 3]), { - headers: { "content-type": "audio/wav" }, - }), + test.each([false, true])( + "applies capture=%s to speech consumed through a cloned response", + async (captureAttachments) => { + initLogger({ + projectName: "groq.test.ts", + projectId: "test-project-id", + captureAttachments, + }); + const wrapped = wrapGroq({ + audio: { + speech: { + create: vi.fn(async (_request: Record) => + Promise.resolve( + new Response(new Uint8Array([1, 2, 3]), { + headers: { "content-type": "audio/wav" }, + }), + ), ), - ), + }, }, - }, - }); + }); - const response = await wrapped.audio.speech.create({ - input: "Hello from Braintrust.", - model: "playai-tts", - voice: "Fritz-PlayAI", - }); - expect(new Uint8Array(await response.clone().arrayBuffer())).toEqual( - new Uint8Array([1, 2, 3]), - ); - expect(response.bodyUsed).toBe(false); - - const spans = (await backgroundLogger.drain()) as Record[]; - const speechSpan = spans.find( - (span) => span.span_attributes?.name === "groq.audio.speech.create", - )!; - const attachment = speechSpan.output.content[0].file - .file_data as Attachment; - expect( - new Uint8Array(await (await attachment.data()).arrayBuffer()), - ).toEqual(new Uint8Array([1, 2, 3])); - }); + const response = await wrapped.audio.speech.create({ + input: "Hello from Braintrust.", + model: "playai-tts", + voice: "Fritz-PlayAI", + }); + expect(new Uint8Array(await response.clone().arrayBuffer())).toEqual( + new Uint8Array([1, 2, 3]), + ); + expect(response.bodyUsed).toBe(false); + + const spans = (await backgroundLogger.drain()) as Record[]; + const speechSpan = spans.find( + (span) => span.span_attributes?.name === "groq.audio.speech.create", + )!; + if (!captureAttachments) { + expect(speechSpan.output.content).toEqual([]); + expect(speechSpan.metrics.time_to_first_token).toEqual( + expect.any(Number), + ); + return; + } + const attachment = speechSpan.output.content[0].file + .file_data as Attachment; + expect( + new Uint8Array(await (await attachment.data()).arrayBuffer()), + ).toEqual(new Uint8Array([1, 2, 3])); + }, + ); }); diff --git a/js/tests/provider-wrappers.test.ts b/js/tests/provider-wrappers.test.ts index 37161415c..ee4ce6363 100644 --- a/js/tests/provider-wrappers.test.ts +++ b/js/tests/provider-wrappers.test.ts @@ -74,7 +74,7 @@ describe("provider wrapper", () => { _exportsForTestingOnly.clearTestBackgroundLogger(); }); - test("OpenAI responses image outputs are converted to attachments when enabled", async () => { + test("OpenAI responses image outputs are converted to attachments", async () => { const originalAutoCaptureAttachments = process.env.BRAINTRUST_CAPTURE_ATTACHMENTS; process.env.BRAINTRUST_CAPTURE_ATTACHMENTS = "true"; @@ -157,7 +157,9 @@ describe("provider wrapper", () => { ).toEqual(mockData); const spans = await backgroundLogger.drain(); expect(spans).toHaveLength(1); - expect(spans[0].output[0].result).toBe(""); + expect(spans[0].output).toEqual([ + { type: "image_generation_call", output_format: "png" }, + ]); } finally { client.responses.create = originalCreate; if (originalAutoCaptureAttachments === undefined) {