diff --git a/.changeset/light-plants-return.md b/.changeset/light-plants-return.md new file mode 100644 index 000000000..93ea58bdb --- /dev/null +++ b/.changeset/light-plants-return.md @@ -0,0 +1,7 @@ +--- +"@braintrust/deepseek-harness": major +"braintrust": major +"@braintrust/openai-agents": minor +--- + +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..cc09b3da9 --- /dev/null +++ b/e2e/helpers/attachment-capture-assertions.ts @@ -0,0 +1,73 @@ +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" | "elevenlabs"; + 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,"); + } + if (options.provider === "elevenlabs") { + expect(payload).toContain("time_to_first_token"); + expect(payload).toContain("annotations"); + expect(payload).not.toContain("audioBase64"); + } else { + 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..f798acf84 --- /dev/null +++ b/e2e/helpers/attachment-capture-scenario.mjs @@ -0,0 +1,149 @@ +import assert from "node:assert/strict"; +import { readFile } from "node:fs/promises"; +import { + initLogger, + wrapOpenAI, + wrapAnthropic, + wrapElevenLabs, +} 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); + }; + } else { + const raw = new sdk.ElevenLabsClient({ + apiKey: process.env.ELEVENLABS_API_KEY, + baseUrl: process.env.ELEVENLABS_BASE_URL, + }); + const client = wrapped ? wrapElevenLabs(raw) : raw; + request = async () => { + const stream = await client.textToSpeech.streamWithTimestamps( + "JBFqnCBsd6RMkjVDRZzb", + { + text: "Streaming timestamps.", + modelId: "eleven_flash_v2_5", + outputFormat: "mp3_44100_128", + }, + ); + let size = 0; + for await (const chunk of stream) size += chunk.audioBase64.length; + assert.ok(size > 0); + }; + } + 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 9ca06b640..2855f4a6c 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, @@ -78,7 +82,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, @@ -100,3 +107,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.capture.mjs b/e2e/scenarios/elevenlabs-instrumentation/scenario.capture.mjs new file mode 100644 index 000000000..b6868e579 --- /dev/null +++ b/e2e/scenarios/elevenlabs-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( + "elevenlabs", + import.meta.url, + await import(process.env.CAPTURE_PACKAGE_NAME), + ), +); diff --git a/e2e/scenarios/elevenlabs-instrumentation/scenario.test.ts b/e2e/scenarios/elevenlabs-instrumentation/scenario.test.ts index e4d82876d..a06fe2767 100644 --- a/e2e/scenarios/elevenlabs-instrumentation/scenario.test.ts +++ b/e2e/scenarios/elevenlabs-instrumentation/scenario.test.ts @@ -1,3 +1,4 @@ +import { defineAttachmentCaptureTests } from "../../helpers/attachment-capture-assertions"; import { describe, expect, test } from "vitest"; import { prepareScenarioDir, @@ -30,6 +31,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-", @@ -143,3 +145,13 @@ for (const variant of ["elevenlabs-v2", "elevenlabs-v2-latest"]) { } }); } + +for (const variant of ["elevenlabs-v2", "elevenlabs-v2-latest"]) { + defineAttachmentCaptureTests({ + scenarioDir, + originalScenarioDir, + provider: "elevenlabs", + packageName: variant.replace("elevenlabs-", "elevenlabs-sdk-"), + variantKey: variant, + }); +} diff --git a/e2e/scenarios/google-genai-instrumentation/scenario.test.ts b/e2e/scenarios/google-genai-instrumentation/scenario.test.ts index 7863a372f..d82e75dd6 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: { + BRAINTRUST_CAPTURE_ATTACHMENTS: "true", + GOOGLE_GENAI_PACKAGE_NAME: scenario.dependencyName, + }, 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: { + BRAINTRUST_CAPTURE_ATTACHMENTS: "true", + GOOGLE_GENAI_PACKAGE_NAME: scenario.dependencyName, + }, 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..0399972b2 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: { + BRAINTRUST_CAPTURE_ATTACHMENTS: "true", + GOOGLE_GENERATIVE_AI_PACKAGE_NAME: scenario.dependencyName, + }, 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: { + BRAINTRUST_CAPTURE_ATTACHMENTS: "true", + GOOGLE_GENERATIVE_AI_PACKAGE_NAME: scenario.dependencyName, + }, 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 f4e0cfa39..683f05e4f 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, @@ -138,3 +139,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/src/index.test.ts b/integrations/deepseek-harness/src/index.test.ts index 6d2c6dcf8..001097b8b 100644 --- a/integrations/deepseek-harness/src/index.test.ts +++ b/integrations/deepseek-harness/src/index.test.ts @@ -45,6 +45,7 @@ const mock = vi.hoisted(() => { } const logger = { + loggingState: { _internalCaptureAttachmentsEnabled: vi.fn(() => true) }, startSpan: (args: Record) => new MockSpan(args), flush: vi.fn(async () => undefined), }; @@ -61,6 +62,9 @@ const mock = vi.hoisted(() => { spans.length = 0; initLogger.mockClear(); logger.flush.mockClear(); + logger.loggingState._internalCaptureAttachmentsEnabled.mockReturnValue( + true, + ); }, spans, }; @@ -228,88 +232,105 @@ 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.logger.loggingState._internalCaptureAttachmentsEnabled.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); + + if (!captureAttachments) { + expect(harness.readImage).not.toHaveBeenCalled(); + expect(mock.attachments).toHaveLength(0); + return; + } + 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]), + ); + + 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, + ], + }, + ], + }); + }, + ); 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..450ff9316 100644 --- a/integrations/deepseek-harness/src/index.ts +++ b/integrations/deepseek-harness/src/index.ts @@ -112,53 +112,56 @@ async function normalizeBlocks( blocks: readonly HarnessContentBlock[], resolveImage: ( ref: HarnessImageAttachmentRef, - ) => Promise, + ) => Promise, ): Promise { if (blocks.length === 0) return ""; if (blocks.every((block) => block.type === "text")) { return contentText(blocks); } - return Promise.all( - blocks.map(async (block) => { - switch (block.type) { - case "text": - return { type: "text", text: block.text ?? "" }; - case "reasoning": - return { type: "reasoning", text: block.text ?? "" }; - case "tool-call": - return { - type: "tool_call", - id: block.id, - name: block.name, - arguments: block.arguments, - }; - case "tool-result": - return { - type: "tool_result", - tool_call_id: block.toolCallId, - content: await normalizeBlocks(block.content ?? [], resolveImage), - ...(block.isError ? { is_error: true } : {}), - }; - case "image": { - const attachment = block.attachment - ? await resolveImage(block.attachment) - : undefined; - return attachment - ? { type: "image_url", image_url: { url: attachment } } - : { type: "image", attachment: block.attachment }; + return ( + await Promise.all( + blocks.map(async (block) => { + switch (block.type) { + case "text": + return { type: "text", text: block.text ?? "" }; + case "reasoning": + return { type: "reasoning", text: block.text ?? "" }; + case "tool-call": + return { + type: "tool_call", + id: block.id, + name: block.name, + arguments: block.arguments, + }; + case "tool-result": + return { + type: "tool_result", + tool_call_id: block.toolCallId, + content: await normalizeBlocks(block.content ?? [], resolveImage), + ...(block.isError ? { is_error: true } : {}), + }; + case "image": { + const attachment = block.attachment + ? await resolveImage(block.attachment) + : undefined; + if (attachment === null) return undefined; + return attachment + ? { type: "image_url", image_url: { url: attachment } } + : { type: "image", attachment: block.attachment }; + } + default: + return { type: block.type }; } - default: - return { type: block.type }; - } - }), - ); + }), + ) + ).filter((block) => block !== undefined); } async function normalizeMessage( message: HarnessMessage, resolveImage: ( ref: HarnessImageAttachmentRef, - ) => Promise, + ) => Promise, ): Promise> { const toolResult = message.content.find( (block) => block.type === "tool-result", @@ -199,7 +202,7 @@ async function normalizeInput( options: HarnessGenerateOptions, resolveImage: ( ref: HarnessImageAttachmentRef, - ) => Promise, + ) => Promise, ): Promise { return [ ...(options.system ? [{ role: "system", content: options.system }] : []), @@ -393,7 +396,9 @@ export function apply(ctx: Context, config: Config = {}): void { const resolveImage = async ( ref: HarnessImageAttachmentRef, - ): Promise => { + ): Promise => { + if (!logger.loggingState._internalCaptureAttachmentsEnabled(logger)) + return null; const attachmentStore = ( ctx as Context & { attachments?: HarnessAttachmentStore } ).attachments; diff --git a/integrations/openai-agents-js/src/attachment-capture.test.ts b/integrations/openai-agents-js/src/attachment-capture.test.ts new file mode 100644 index 000000000..80e60d93b --- /dev/null +++ b/integrations/openai-agents-js/src/attachment-capture.test.ts @@ -0,0 +1,68 @@ +import { afterEach, beforeAll, beforeEach, expect, it, vi } from "vitest"; +import { _exportsForTestingOnly, initLogger } from "braintrust"; +import { OpenAIAgentsTraceProcessor } from "./index"; +import type { AgentsSpan, AgentsTrace } from "./types"; + +let background: ReturnType< + typeof _exportsForTestingOnly.useTestBackgroundLogger +>; +beforeAll(() => _exportsForTestingOnly.simulateLoginForTests()); +beforeEach(() => { + background = _exportsForTestingOnly.useTestBackgroundLogger(); +}); +afterEach(() => { + _exportsForTestingOnly.clearTestBackgroundLogger(); + vi.unstubAllEnvs(); +}); + +it.each([true, false])( + "uses its local logger policy (%s) when trace events arrive outside the original context", + async (captureAttachments) => { + vi.stubEnv("BRAINTRUST_CAPTURE_ATTACHMENTS", String(!captureAttachments)); + const options = { + projectId: "test-project-id", + projectName: "tmp-luca-agents-capture", + }; + const logger = initLogger({ + ...options, + captureAttachments, + setCurrent: false, + }); + const processor = new OpenAIAgentsTraceProcessor({ logger }); + const trace: AgentsTrace = { + type: "trace", + traceId: "trace", + name: "test trace", + groupId: null, + }; + const span: AgentsSpan = { + type: "trace.span", + traceId: "trace", + spanId: "response", + parentId: null, + startedAt: null, + endedAt: null, + error: null, + spanData: { + type: "response", + _input: [{ type: "input_image", image: "data:image/png;base64,AQID" }], + _response: { + output: [{ type: "image_generation_call", result: "AQID" }], + }, + }, + }; + await processor.onTraceStart(trace); + await processor.onSpanStart(span); + initLogger({ ...options, captureAttachments: !captureAttachments }); + await processor.onSpanEnd(span); + await processor.onTraceEnd(trace); + const payload = JSON.stringify(await background.drain()); + if (captureAttachments) expect(payload).toContain("braintrust_attachment"); + else { + expect(payload).not.toContain("input_image"); + expect(payload).not.toContain("image_generation_call"); + expect(payload).not.toContain("AQID"); + expect(payload).not.toContain("braintrust_attachment"); + } + }, +); diff --git a/integrations/openai-agents-js/src/index.ts b/integrations/openai-agents-js/src/index.ts index e40b95173..cbd81c10f 100644 --- a/integrations/openai-agents-js/src/index.ts +++ b/integrations/openai-agents-js/src/index.ts @@ -7,6 +7,9 @@ import { currentSpan, NOOP_SPAN, Attachment, + _internalGetGlobalState, + getSpanParentObject, + withCurrent, } from "braintrust"; import { SpanType, @@ -143,12 +146,25 @@ export class OpenAIAgentsTraceProcessor { private processInputImages(input: any): any { if (Array.isArray(input)) { - return input.map((item) => this.processInputImages(item)); + return input + .map((item) => this.processInputImages(item)) + .filter((item) => item !== undefined); } if (input && typeof input === "object") { // Handle input_image type with base64 image data if (input.type === "input_image" && typeof input.image === "string") { + if (/^https?:\/\//i.test(input.image)) return input; + if ( + !_internalGetGlobalState()._internalCaptureAttachmentsEnabled( + getSpanParentObject(), + ) + ) { + const { image: _image, ...metadata } = input; + return Object.keys(metadata).some((key) => key !== "type") + ? metadata + : undefined; + } let imageData = input.image; // Strip data URI prefix if present (e.g., "data:image/png;base64,") @@ -193,7 +209,8 @@ export class OpenAIAgentsTraceProcessor { // Recursively process nested objects const result: any = {}; for (const [key, value] of Object.entries(input)) { - result[key] = this.processInputImages(value); + const processed = this.processInputImages(value); + if (processed !== undefined) result[key] = processed; } return result; } @@ -203,12 +220,24 @@ export class OpenAIAgentsTraceProcessor { private processOutputImages(output: any): any { if (Array.isArray(output)) { - return output.map((item) => this.processOutputImages(item)); + return output + .map((item) => this.processOutputImages(item)) + .filter((item) => item !== undefined); } if (output && typeof output === "object") { // Handle image_generation_call type - convert result to attachment if (output.type === "image_generation_call" && output.result) { + if ( + !_internalGetGlobalState()._internalCaptureAttachmentsEnabled( + getSpanParentObject(), + ) + ) { + const { result: _result, ...metadata } = output; + return Object.keys(metadata).some((key) => key !== "type") + ? metadata + : undefined; + } let resultData = output.result; // Use output_format from the response @@ -256,7 +285,8 @@ export class OpenAIAgentsTraceProcessor { // Recursively process nested objects const result: any = {}; for (const [key, value] of Object.entries(output)) { - result[key] = this.processOutputImages(value); + const processed = this.processOutputImages(value); + if (processed !== undefined) result[key] = processed; } return result; } @@ -635,7 +665,9 @@ export class OpenAIAgentsTraceProcessor { const braintrustSpan = traceData.childSpans.get(span.spanId); if (braintrustSpan) { - const logData = this.extractLogData(span); + const logData = withCurrent(braintrustSpan, () => + this.extractLogData(span), + ); braintrustSpan.log({ error: span.error, ...logData, diff --git a/integrations/openai-agents-js/src/openai-agents-integration.test.ts b/integrations/openai-agents-js/src/openai-agents-integration.test.ts index 317d04b7e..7803a453f 100644 --- a/integrations/openai-agents-js/src/openai-agents-integration.test.ts +++ b/integrations/openai-agents-js/src/openai-agents-integration.test.ts @@ -85,6 +85,7 @@ describe( beforeEach(() => { backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); _logger = initLogger({ + captureAttachments: true, projectName: "openai-agents.test.ts", projectId: "test-openai-agents", }); diff --git a/js/src/attachment-capture.test.ts b/js/src/attachment-capture.test.ts new file mode 100644 index 000000000..3150ea6fe --- /dev/null +++ b/js/src/attachment-capture.test.ts @@ -0,0 +1,371 @@ +import { afterEach, beforeAll, beforeEach, expect, it, vi } from "vitest"; +import { + Attachment, + BraintrustState, + CAPTURE_ATTACHMENTS, + startSpan, + _exportsForTestingOnly, + _internalGetGlobalState, + initLogger, + withCurrent, +} from "./logger"; +import { configureNode } from "./node/config"; +import iso from "./isomorph"; +import * as byteUtils from "../util/index"; +import { + isAutoCaptureAttachmentsEnabled, + processInputAttachments, +} from "./wrappers/attachment-utils"; +import { processAttachmentsInInput } from "./instrumentation/plugins/anthropic-plugin"; +import { processAISDKGenerateImageOutput } from "./instrumentation/plugins/ai-sdk-plugin"; +import { extractOllamaChatInput } from "./instrumentation/plugins/ollama-plugin"; +import { openAIChannels } from "./instrumentation/plugins/openai-channels"; +import { elevenLabsChannels } from "./instrumentation/plugins/elevenlabs-channels"; +import { groqChannels } from "./instrumentation/plugins/groq-channels"; +import { googleGenAIChannels } from "./instrumentation/plugins/google-genai-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([true, false])( + "explicit %s overrides an opposing environment setting", + (captureAttachments) => { + vi.stubEnv("BRAINTRUST_CAPTURE_ATTACHMENTS", String(!captureAttachments)); + const logger = initLogger({ ...loggerOptions, captureAttachments }); + expect(isAutoCaptureAttachmentsEnabled(logger)).toBe(captureAttachments); + expect(isAutoCaptureAttachmentsEnabled()).toBe(captureAttachments); + initLogger(loggerOptions); + expect(isAutoCaptureAttachmentsEnabled()).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(isAutoCaptureAttachmentsEnabled(local)).toBe(false); + initLogger({ + ...loggerOptions, + setCurrent: false, + captureAttachments: false, + }); + expect(isAutoCaptureAttachmentsEnabled()).toBe(true); + vi.stubEnv("BRAINTRUST_CAPTURE_ATTACHMENTS", "true"); + initLogger({ ...loggerOptions, captureAttachments: false }); + expect(isAutoCaptureAttachmentsEnabled(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(isAutoCaptureAttachmentsEnabled(otherLogger)).toBe(false); + expect(other._internalCaptureAttachmentsEnabled()).toBe(false); + expect(isAutoCaptureAttachmentsEnabled()).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.each([true, false])( + "retains a local policy (%s) when consuming an OpenAI stream outside its trace", + async (captureAttachments) => { + const local = initLogger({ + ...loggerOptions, + setCurrent: false, + captureAttachments, + }); + async function* events() { + yield { + type: "image_generation.completed", + b64_json: "AQID", + output_format: "png", + }; + } + const stream = await local.traced(() => + openAIChannels.imagesGenerate.invoke( + async () => events(), + undefined, + [{ prompt: "draw", stream: true }], + {}, + ), + ); + initLogger({ ...loggerOptions, captureAttachments: !captureAttachments }); + for await (const event of stream as AsyncIterable) + expect(event).toBeDefined(); + const rows = (await background.drain()) as Array<{ + input?: unknown; + output?: unknown; + metrics?: Record; + }>; + const output = rows.find((row) => row.output)?.output as { + content: { image_url: { url: unknown } }[]; + }; + if (captureAttachments) + expect(output.content[0].image_url.url).toBeInstanceOf(Attachment); + else expect(output.content).toEqual([]); + }, +); + +it("uses the owning span for delayed Google media output", async () => { + const local = initLogger({ + ...loggerOptions, + setCurrent: false, + captureAttachments: false, + }); + let resolve!: (value: { + generatedImages: { image: { imageBytes: string; mimeType: string } }[]; + }) => void; + const promise = new Promise<{ + generatedImages: { image: { imageBytes: string; mimeType: string } }[]; + }>((done) => { + resolve = done; + }); + const pending = local.traced(() => + googleGenAIChannels.generateImages.invoke( + () => promise, + undefined, + [{ model: "model", prompt: "draw" }], + {}, + ), + ); + initLogger({ ...loggerOptions, captureAttachments: true }); + const binary = vi.fn((): string => { + throw new Error("must not read skipped media"); + }); + resolve({ + generatedImages: [ + { + image: { + get imageBytes() { + return binary(); + }, + mimeType: "image/png", + }, + }, + ], + }); + await pending; + expect(binary).not.toHaveBeenCalled(); + const rows = (await background.drain()) as Array<{ + input?: unknown; + output?: unknown; + metrics?: Record; + }>; + expect(JSON.stringify(rows)).not.toContain("AQID"); + expect(JSON.stringify(rows)).not.toContain("file_data"); +}); + +it("keeps ElevenLabs transcripts and timing without decoding or retaining audio", async () => { + const local = initLogger({ + ...loggerOptions, + setCurrent: false, + captureAttachments: false, + }); + const decode = vi.spyOn(globalThis, "atob"); + async function* audio() { + yield { audioBase64: "AQID", alignment: { characters: ["h", "i"] } }; + } + const span = local.startSpan(); + const stream = await withCurrent(span, () => + elevenLabsChannels.streamWithTimestamps.invoke( + async () => audio(), + undefined, + ["voice", { text: "hi" }], + {}, + ), + ); + initLogger({ ...loggerOptions, captureAttachments: true }); + for await (const chunk of stream) expect(chunk.audioBase64).toBe("AQID"); + span.end(); + const rows = (await background.drain()) as Array<{ + input?: unknown; + output?: unknown; + metrics?: Record; + }>; + expect(decode).not.toHaveBeenCalled(); + expect(JSON.stringify(rows)).not.toContain("AQID"); + expect(JSON.stringify(rows)).not.toContain("file_data"); + expect( + rows.some((row) => row.metrics?.time_to_first_token !== undefined), + ).toBe(true); + expect(rows.find((row) => row.output)?.output).toMatchObject({ + content: [], + annotations: [{ alignment: { characters: ["h", "i"] } }], + }); +}); + +it("does not read Groq speech blobs or retain audio when disabled", async () => { + const local = initLogger({ + ...loggerOptions, + setCurrent: false, + captureAttachments: false, + }); + const response = new Response(new Uint8Array([1, 2, 3]), { + headers: { "content-type": "audio/wav" }, + }); + const blobRead = vi.spyOn(Blob.prototype, "arrayBuffer"); + const concatenate = vi.spyOn(byteUtils, "concatUint8Arrays"); + const result = await local.traced(() => + groqChannels.audioSpeechCreate.tracePromise(async () => response, { + arguments: [{ model: "model", input: "hi", voice: "voice" }], + }), + ); + initLogger({ ...loggerOptions, captureAttachments: true }); + const blob = await result.blob(); + expect(blob.size).toBe(3); + expect(blobRead).not.toHaveBeenCalled(); + expect(concatenate).not.toHaveBeenCalled(); + const rows = (await background.drain()) as Array<{ output?: unknown }>; + expect(JSON.stringify(rows)).not.toContain("braintrust_attachment"); + expect(rows.find((row) => row.output)?.output).toEqual({ content: [] }); +}); + +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(isAutoCaptureAttachmentsEnabled(deferred)).toBe(false); + expect(isAutoCaptureAttachmentsEnabled(child)).toBe(false); + child.end(); + deferred.end(); +}); diff --git a/js/src/instrumentation/core/channel-tracing.ts b/js/src/instrumentation/core/channel-tracing.ts index a8ecc819c..44d684684 100644 --- a/js/src/instrumentation/core/channel-tracing.ts +++ b/js/src/instrumentation/core/channel-tracing.ts @@ -4,6 +4,7 @@ import { _internalGetGlobalState, BRAINTRUST_CURRENT_SPAN_STORE, startSpan, + withCurrent, } from "../../logger"; import type { CurrentSpanStore, Span } from "../../logger"; import { @@ -219,10 +220,8 @@ function startSpanForEvent< const startTime = getCurrentUnixTimestamp(); try { - const { input, metadata } = config.extractInput( - event.arguments, - event as StartOf, - span, + const { input, metadata } = withCurrent(span, () => + config.extractInput(event.arguments, event as StartOf, span), ); span.log({ input, @@ -480,9 +479,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, @@ -611,19 +609,24 @@ export function traceStreamingChannel( let metadata: Record | 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; } 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, @@ -719,9 +722,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/plugins/ai-sdk-plugin.test.ts b/js/src/instrumentation/plugins/ai-sdk-plugin.test.ts index 015b3763e..5214927d1 100644 --- a/js/src/instrumentation/plugins/ai-sdk-plugin.test.ts +++ b/js/src/instrumentation/plugins/ai-sdk-plugin.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/plugins/ai-sdk-plugin.ts b/js/src/instrumentation/plugins/ai-sdk-plugin.ts index 989216a67..1ad6b083f 100644 --- a/js/src/instrumentation/plugins/ai-sdk-plugin.ts +++ b/js/src/instrumentation/plugins/ai-sdk-plugin.ts @@ -30,6 +30,8 @@ import { import { convertDataToBlob, getExtensionFromMediaType, + isAutoCaptureAttachmentsEnabled, + omitMediaData, } from "../../wrappers/attachment-utils"; import { normalizeAISDKLoggedOutput } from "../../wrappers/ai-sdk/normalize-logged-output"; import { serializeAISDKToolsForLogging } from "../../wrappers/ai-sdk/tool-serialization"; @@ -1570,7 +1572,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), }; } @@ -1586,14 +1590,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 { @@ -1615,6 +1623,7 @@ const processContentPart = (part: any): any => { part.image, part.mimeType || part.mediaType, ); + if (imageAttachment === undefined) return omitMediaData(part, "image"); if (imageAttachment) { return { ...part, @@ -1633,6 +1642,7 @@ const processContentPart = (part: any): any => { part.mimeType || part.mediaType, part.name || part.filename, ); + if (fileAttachment === undefined) return omitMediaData(part, "data"); if (fileAttachment) { return { ...part, @@ -1644,6 +1654,11 @@ const processContentPart = (part: any): any => { if (part.type === "image_url" && part.image_url) { if (typeof part.image_url === "object" && part.image_url.url) { const imageAttachment = convertImageToAttachment(part.image_url.url); + if (imageAttachment === undefined) + return omitMediaData({ + ...part, + image_url: omitMediaData(part.image_url, "url"), + }); if (imageAttachment) { return { ...part, @@ -1666,7 +1681,15 @@ const processContentPart = (part: any): any => { const convertImageToAttachment = ( image: any, explicitMimeType?: string, -): Attachment | null => { +): Attachment | undefined | null => { + if (!isAutoCaptureAttachmentsEnabled()) { + return image instanceof Attachment + ? image + : image instanceof URL || + (typeof image === "string" && /^https?:/.test(image)) + ? null + : undefined; + } try { if (typeof image === "string" && image.startsWith("data:")) { const [mimeTypeSection, base64Data] = image.split(","); @@ -1717,7 +1740,15 @@ const convertDataToAttachment = ( data: any, mimeType: string, filename?: string, -): Attachment | null => { +): Attachment | undefined | null => { + if (!isAutoCaptureAttachmentsEnabled()) { + return data instanceof Attachment + ? data + : data instanceof URL || + (typeof data === "string" && /^https?:/.test(data)) + ? null + : undefined; + } if (!mimeType) return null; try { @@ -4030,9 +4061,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; @@ -4046,6 +4078,8 @@ function convertAISDKGeneratedFileToAttachment( return file; } + if (!isAutoCaptureAttachmentsEnabled()) return undefined; + const generatedFile = file as AISDKGeneratedFile & Record; const generatedMediaType = safeSerializableFieldRead( generatedFile, diff --git a/js/src/instrumentation/plugins/anthropic-plugin.test.ts b/js/src/instrumentation/plugins/anthropic-plugin.test.ts index a711c3b7a..4f1f22e7a 100644 --- a/js/src/instrumentation/plugins/anthropic-plugin.test.ts +++ b/js/src/instrumentation/plugins/anthropic-plugin.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, + ), }, })); @@ -22,6 +25,7 @@ const aggregateAnthropicStreamChunksForTest = (chunks: unknown[]) => // Mock startSpan from logger vi.mock("../../logger", () => ({ + withCurrent: (_span: unknown, callback: () => unknown) => callback(), startSpan: vi.fn(() => ({ log: vi.fn(), end: vi.fn(), diff --git a/js/src/instrumentation/plugins/anthropic-plugin.ts b/js/src/instrumentation/plugins/anthropic-plugin.ts index e184c927d..b4e5ca367 100644 --- a/js/src/instrumentation/plugins/anthropic-plugin.ts +++ b/js/src/instrumentation/plugins/anthropic-plugin.ts @@ -1,3 +1,7 @@ +import { + isAutoCaptureAttachmentsEnabled, + omitMediaData, +} from "../../wrappers/attachment-utils"; import { BasePlugin, toLoggedError } from "../core"; import { traceStreamingChannel, unsubscribeAll } from "../core/channel-tracing"; import { isAsyncIterable, patchStreamIfNeeded } from "../core/stream-patcher"; @@ -1651,15 +1655,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 = isAutoCaptureAttachmentsEnabled(), +): 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), @@ -1669,7 +1683,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/plugins/elevenlabs-plugin.test.ts b/js/src/instrumentation/plugins/elevenlabs-plugin.test.ts index 0334d7638..ea55eaf97 100644 --- a/js/src/instrumentation/plugins/elevenlabs-plugin.test.ts +++ b/js/src/instrumentation/plugins/elevenlabs-plugin.test.ts @@ -1,4 +1,12 @@ -import { afterEach, beforeAll, beforeEach, describe, expect, it } from "vitest"; +import { + afterEach, + beforeAll, + beforeEach, + describe, + expect, + it, + vi, +} from "vitest"; import { Readable } from "node:stream"; import { ReadableStream, @@ -11,6 +19,8 @@ import { elevenLabsChannels } from "./elevenlabs-channels"; import { wrapElevenLabs } from "../../wrappers/elevenlabs"; configureNode(); +beforeEach(() => vi.stubEnv("BRAINTRUST_CAPTURE_ATTACHMENTS", "true")); +afterEach(() => vi.unstubAllEnvs()); const request = { text: "Hello", modelId: "eleven_flash_v2_5" }; describe("ElevenLabs instrumentation", () => { diff --git a/js/src/instrumentation/plugins/elevenlabs-plugin.ts b/js/src/instrumentation/plugins/elevenlabs-plugin.ts index fc3abed71..393d620da 100644 --- a/js/src/instrumentation/plugins/elevenlabs-plugin.ts +++ b/js/src/instrumentation/plugins/elevenlabs-plugin.ts @@ -12,7 +12,10 @@ import type { ElevenLabsTranscription, ElevenLabsTranscriptionRequest, } from "../../vendor-sdk-types/elevenlabs"; -import { getExtensionFromMediaType } from "../../wrappers/attachment-utils"; +import { + getExtensionFromMediaType, + isAutoCaptureAttachmentsEnabled, +} from "../../wrappers/attachment-utils"; import { isAutoInstrumentationSuppressed, runWithAutoInstrumentationSuppressed, @@ -84,8 +87,10 @@ export class ElevenLabsPlugin extends BasePlugin { file instanceof Blob ? file.type || "application/octet-stream" : "application/octet-stream"; - const blob = - file instanceof Blob + const captureAttachments = isAutoCaptureAttachmentsEnabled(); + const blob = !captureAttachments + ? undefined + : file instanceof Blob ? file : file instanceof Uint8Array ? new Blob([new Uint8Array(file)], { type: contentType }) @@ -135,7 +140,8 @@ export class ElevenLabsPlugin extends BasePlugin { }, ([request], span) => { const file = request.file; - if (!isAsyncIterable(file)) return; + if (!isAutoCaptureAttachmentsEnabled(span) || !isAsyncIterable(file)) + return; const chunks: Uint8Array[] = []; const filename = isObject(file) && typeof file.path === "string" @@ -302,31 +308,8 @@ function captureSpeech( started: number, headers?: Headers, ): void { - const format = request.outputFormat ?? "mp3_44100_128"; - const formatType = format.startsWith("mp3_") - ? "audio/mpeg" - : format.startsWith("pcm_") - ? "audio/pcm" - : format.startsWith("opus_") - ? "audio/ogg" - : format.startsWith("ulaw_") - ? "audio/basic" - : format.startsWith("alaw_") - ? "audio/x-alaw" - : undefined; - const headerType = headers?.get("content-type")?.split(";")[0]; - const contentType = headerType?.startsWith("audio/") - ? headerType - : formatType; - const disposition = headers?.get("content-disposition"); - const headerFilename = disposition?.match( - /filename="([^"]+)"|filename=([^;]+)/i, - ); - const filename = - headerFilename?.[1] ?? - headerFilename?.[2]?.trim() ?? - `speech.${contentType ? getExtensionFromMediaType(contentType) : "bin"}`; - const chunks: Uint8Array[] = []; + const captureAttachments = isAutoCaptureAttachmentsEnabled(span); + const chunks: Uint8Array[] | undefined = captureAttachments ? [] : undefined; const alignments: unknown[] = []; let bytes = 0; let first = true; @@ -339,7 +322,7 @@ function captureSpeech( }); } first = false; - chunks.push(new Uint8Array(chunk)); + if (chunks) chunks.push(new Uint8Array(chunk)); bytes += chunk.byteLength; }; const complete = () => { @@ -347,25 +330,51 @@ function captureSpeech( stopped = true; try { const content = []; - if (contentType && bytes) { - const data = new Uint8Array(bytes); - let offset = 0; - for (const chunk of chunks) { - data.set(chunk, offset); - offset += chunk.length; - } - content.push({ - type: "file", - file: { - filename, - byte_size: bytes, - file_data: new Attachment({ - data: new Blob([data], { type: contentType }), + if (chunks && bytes) { + const format = request.outputFormat ?? "mp3_44100_128"; + const formatType = format.startsWith("mp3_") + ? "audio/mpeg" + : format.startsWith("pcm_") + ? "audio/pcm" + : format.startsWith("opus_") + ? "audio/ogg" + : format.startsWith("ulaw_") + ? "audio/basic" + : format.startsWith("alaw_") + ? "audio/x-alaw" + : undefined; + const headerType = headers?.get("content-type")?.split(";")[0]; + const contentType = headerType?.startsWith("audio/") + ? headerType + : formatType; + const disposition = headers?.get("content-disposition"); + const headerFilename = disposition?.match( + /filename="([^"]+)"|filename=([^;]+)/i, + ); + const filename = + headerFilename?.[1] ?? + headerFilename?.[2]?.trim() ?? + `speech.${contentType ? getExtensionFromMediaType(contentType) : "bin"}`; + if (contentType) { + const data = new Uint8Array(bytes); + let offset = 0; + for (const chunk of chunks) { + data.set(chunk, offset); + offset += chunk.length; + } + content.push({ + type: "file", + file: { filename, - contentType, - }), - }, - }); + byte_size: bytes, + file_data: new Attachment({ + data: new Blob([data], { type: contentType }), + filename, + contentType, + }), + }, + }); + } } span.log({ output: { @@ -374,18 +383,26 @@ function captureSpeech( }, }); } finally { - chunks.length = 0; + if (chunks) chunks.length = 0; finish(); } }; const cancel: Finish = (error) => { stopped = true; - chunks.length = 0; + if (chunks) chunks.length = 0; 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) { + first = false; + if (method.startsWith("stream")) + span.log({ + metrics: { time_to_first_token: Date.now() / 1000 - started }, + }); + } if (chunk.alignment || chunk.normalizedAlignment) alignments.push({ alignment: chunk.alignment, diff --git a/js/src/instrumentation/plugins/google-genai-plugin.test.ts b/js/src/instrumentation/plugins/google-genai-plugin.test.ts index 34f8ac3da..7b8682be1 100644 --- a/js/src/instrumentation/plugins/google-genai-plugin.test.ts +++ b/js/src/instrumentation/plugins/google-genai-plugin.test.ts @@ -19,6 +19,9 @@ vi.mock("../../isomorph", () => ({ }; }), newTracingChannel: vi.fn(), + getEnv: vi.fn((name: string) => + name === "BRAINTRUST_CAPTURE_ATTACHMENTS" ? "true" : undefined, + ), }, })); @@ -31,6 +34,8 @@ const mockStartSpan = vi.mocked(startSpan); // Mock logger vi.mock("../../logger", () => ({ + BaseAttachment: class {}, + CAPTURE_ATTACHMENTS: Symbol.for("braintrust.captureAttachments"), startSpan: vi.fn(() => ({ log: vi.fn(), end: vi.fn(), diff --git a/js/src/instrumentation/plugins/google-genai-plugin.ts b/js/src/instrumentation/plugins/google-genai-plugin.ts index b4ec2b7ee..a61d6e4b1 100644 --- a/js/src/instrumentation/plugins/google-genai-plugin.ts +++ b/js/src/instrumentation/plugins/google-genai-plugin.ts @@ -1,6 +1,8 @@ import { uint8ArrayToBase64 } from "../../../util/bytes"; import { getExtensionFromMediaType, + isAutoCaptureAttachmentsEnabled, + omitMediaData, processInputAttachments, } from "../../wrappers/attachment-utils"; import { debugLogger } from "../../debug-logger"; @@ -15,6 +17,7 @@ import type { IsoChannelHandlers, IsoTracingChannel } from "../../isomorph"; import { _internalGetGlobalState, Attachment, + CAPTURE_ATTACHMENTS, currentSpan, BRAINTRUST_CURRENT_SPAN_STORE, startSpan as startBaseSpan, @@ -69,6 +72,7 @@ type GenerateContentStreamEvent = googleGenAIInput?: Record; googleGenAIMetadata?: Record; googleGenAIStartTime?: number; + captureAttachments?: boolean; }; type SpanState = { @@ -205,7 +209,9 @@ export class GoogleGenAIPlugin extends BasePlugin { spanState.startTime, ), ), - output: serializeGenerateContentOutput(event.result), + output: withCurrent(spanState.span, () => + serializeGenerateContentOutput(event.result), + ), }); } finally { spanState.span.end(); @@ -240,10 +246,12 @@ export class GoogleGenAIPlugin extends BasePlugin { streamEvent.googleGenAIMetadata = extractGenerateContentMetadata(params); streamEvent.googleGenAIStartTime = getCurrentUnixTimestamp(); + streamEvent.captureAttachments = isAutoCaptureAttachmentsEnabled(); }, asyncEnd: (event) => { const streamEvent = event as GenerateContentStreamEvent; patchGoogleGenAIStreamingResult({ + captureAttachments: streamEvent.captureAttachments, input: streamEvent.googleGenAIInput, metadata: streamEvent.googleGenAIMetadata, startTime: streamEvent.googleGenAIStartTime, @@ -543,7 +551,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:`, @@ -640,12 +650,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 = isAutoCaptureAttachmentsEnabled(), + } = args; if ( !input || @@ -670,6 +687,7 @@ function patchGoogleGenAIStreamingResult(args: { withSpanInstrumentationName( { name: "generate_content_stream", + [CAPTURE_ATTACHMENTS]: captureAttachments, spanAttributes: { type: SpanTypeAttribute.LLM, }, @@ -791,7 +809,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) { @@ -800,6 +846,7 @@ function patchGoogleGenAIStreamingResult(args: { chunks, requestStartTime, firstTokenTime, + captureAttachments, ), }); } @@ -827,6 +874,7 @@ function patchGoogleGenAIStreamingResult(args: { chunks, requestStartTime, firstTokenTime, + captureAttachments, ), }); } else { @@ -902,7 +950,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), }, } : {}), @@ -1125,12 +1175,19 @@ function serializeGoogleGenAIImage( fallbackMimeType?: string, purpose?: "input" | "reference" | "mask", ): Record | undefined { + const captureAttachments = isAutoCaptureAttachmentsEnabled(); + if (!captureAttachments && !image.gcsUri) return undefined; const mimeType = image.mimeType ?? fallbackMimeType ?? "image/png"; const filename = `${filenameStem}.${getExtensionFromMediaType(mimeType)}`; const media = image.gcsUri ?? (image.imageBytes - ? createAttachmentFromInlineData(image.imageBytes, mimeType, filename) + ? createAttachmentFromInlineData( + image.imageBytes, + mimeType, + filename, + captureAttachments, + ) : undefined); if (!media) { return undefined; @@ -1146,12 +1203,19 @@ function serializeGoogleGenAIVideo( video: GoogleGenAIVideo, filenameStem: string, ): Record | undefined { + const captureAttachments = isAutoCaptureAttachmentsEnabled(); + if (!captureAttachments && !video.uri) return undefined; const mimeType = video.mimeType ?? "video/mp4"; const filename = `${filenameStem}.${getExtensionFromMediaType(mimeType)}`; const media = video.uri ?? (video.videoBytes - ? createAttachmentFromInlineData(video.videoBytes, mimeType, filename) + ? createAttachmentFromInlineData( + video.videoBytes, + mimeType, + filename, + captureAttachments, + ) : undefined); return media ? { type: "file", file: { filename, file_data: media } } @@ -1200,6 +1264,7 @@ function serializeEmbedContentInput( if (part.text !== undefined) return [{ type: "text", text: part.text }]; const media = part.inlineData ?? part.fileData; if (!media) return []; + if ("data" in media && !isAutoCaptureAttachmentsEnabled()) return []; const data = "data" in media ? `data:${media.mimeType};base64,${typeof media.data === "string" ? media.data : uint8ArrayToBase64(media.data)}` @@ -1346,7 +1411,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; @@ -1362,14 +1429,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 = isAutoCaptureAttachmentsEnabled(), +): 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/") @@ -1377,7 +1457,10 @@ function serializePart(part: GoogleGenAIPart): unknown { : { file: { file_data: attachment, - filename: attachment.reference.filename, + filename: + attachment instanceof Attachment + ? attachment.reference.filename + : `file.${getExtensionFromMediaType(mimeType)}`, }, }; } @@ -1594,12 +1677,18 @@ function serializeInteractionValue( : "mimeType" in dict && typeof dict.mimeType === "string" ? dict.mimeType : undefined; + const captureAttachments = isAutoCaptureAttachmentsEnabled(); 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" && mimeType && !captureAttachments) continue; + const entry: unknown = Reflect.get(dict, key); if (key === "data" && attachment) { serialized[key] = attachment; } else { @@ -1616,8 +1705,10 @@ function serializeInteractionValue( function createAttachmentFromInlineData( data: unknown, mimeType?: string, - filename = `file.${mimeType ? getExtensionFromMediaType(mimeType) : "bin"}`, + filename?: string, + captureAttachments = isAutoCaptureAttachmentsEnabled(), ): Attachment | null { + if (!captureAttachments) return null; if ( !( data instanceof Uint8Array || @@ -1650,7 +1741,9 @@ function createAttachmentFromInlineData( return new Attachment({ data: arrayBuffer, - filename, + filename: + filename ?? + `file.${mimeType ? getExtensionFromMediaType(mimeType) : "bin"}`, contentType: mimeType || "application/octet-stream", }); } @@ -1968,6 +2061,7 @@ function aggregateGenerateContentChunks( chunks: GoogleGenAIGenerateContentResponse[], startTime: number, firstTokenTime: number | null, + captureAttachments: boolean, ): { aggregated: Record; metrics: Record; @@ -2023,7 +2117,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); } diff --git a/js/src/instrumentation/plugins/google-generative-ai-plugin.test.ts b/js/src/instrumentation/plugins/google-generative-ai-plugin.test.ts index b73680cc3..4ff58257a 100644 --- a/js/src/instrumentation/plugins/google-generative-ai-plugin.test.ts +++ b/js/src/instrumentation/plugins/google-generative-ai-plugin.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,36 @@ 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([true, false])( + "processes inline inputs with capture=%s without mutating the provider request", + async (captureAttachments) => { + initLogger({ + projectId: "test-project-id", + projectName: "tmp-luca-google-generative-ai-tests", + 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]).not.toHaveProperty( + "input.contents.0.parts.0.inlineData.data", + ); + 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/plugins/google-generative-ai-plugin.ts b/js/src/instrumentation/plugins/google-generative-ai-plugin.ts index d013bf031..50ba566e0 100644 --- a/js/src/instrumentation/plugins/google-generative-ai-plugin.ts +++ b/js/src/instrumentation/plugins/google-generative-ai-plugin.ts @@ -6,7 +6,11 @@ import { withSpanInstrumentationName, } from "../../span-origin"; import { getCurrentUnixTimestamp, isObject } from "../../util"; -import { processInputAttachments } from "../../wrappers/attachment-utils"; +import { + isAutoCaptureAttachmentsEnabled, + omitMediaData, + processInputAttachments, +} from "../../wrappers/attachment-utils"; import type { GenerativeAIChat, GenerativeAIConfig, @@ -182,9 +186,10 @@ function interceptCall< number, NonNullable[number] >(); + const captureAttachments = isAutoCaptureAttachmentsEnabled(span); patchStreamIfNeeded(value.stream, { aroundNext: (callback) => withCurrent(span, callback), - onChunk: (chunk) => { + shouldCollect: (chunk) => { for (const candidate of chunk.candidates ?? []) { const index = candidate.index ?? 0; const previous = candidates.get(index); @@ -196,7 +201,12 @@ function interceptCall< ...last, text: last.text + part.text, }; - } else parts.push(part); + } else + parts.push( + !captureAttachments && part.inlineData + ? convertAttachments(part, false) + : part, + ); } candidates.set(index, { ...previous, @@ -230,6 +240,7 @@ function interceptCall< }, }); } + return false; }, onComplete: () => {}, onError: (error) => @@ -393,16 +404,28 @@ 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 { +function convertAttachments( + value: T, + captureAttachments = isAutoCaptureAttachmentsEnabled(), +): T { 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: omitMediaData(node.inlineData, "data"), + }); const { data, mimeType } = node.inlineData; - const processed = processInputAttachments({ - type: "file", - file: { file_data: `data:${mimeType};base64,${data}` }, - }); + const processed = processInputAttachments( + { + type: "file", + file: { file_data: `data:${mimeType};base64,${data}` }, + }, + captureAttachments, + ); if (typeof processed.file.file_data === "string") throw new Error("Unable to convert Google inline data"); return { @@ -445,10 +468,13 @@ function logResponse(span: Span, response: GenerativeAIResponse): void { metrics.completion_reasoning_tokens = usage.thoughtsTokenCount; } span.log({ - output: convertAttachments({ - candidates: response.candidates, - promptFeedback: response.promptFeedback, - }), + output: convertAttachments( + { + candidates: response.candidates, + promptFeedback: response.promptFeedback, + }, + isAutoCaptureAttachmentsEnabled(span), + ), metrics, ...(response.modelVersion ? { metadata: { model: response.modelVersion } } diff --git a/js/src/instrumentation/plugins/groq-plugin.ts b/js/src/instrumentation/plugins/groq-plugin.ts index b37623a46..64a9b8de3 100644 --- a/js/src/instrumentation/plugins/groq-plugin.ts +++ b/js/src/instrumentation/plugins/groq-plugin.ts @@ -13,6 +13,7 @@ import { Attachment, withCurrent, type Span } from "../../logger"; import { convertDataToBlob, getExtensionFromMediaType, + isAutoCaptureAttachmentsEnabled, processInputAttachments, } from "../../wrappers/attachment-utils"; import { getCurrentUnixTimestamp } from "../../util"; @@ -94,7 +95,9 @@ export class GroqPlugin extends BasePlugin { }, metadata: { model: params.model, provider: "groq" }, }), - extractOutput: () => ({ content: [] }), + extractOutput: () => ({ + content: [], + }), extractMetrics: () => ({}), patchResult: ({ endEvent, result, span, startTime }) => captureGroqSpeechResponse( @@ -144,7 +147,8 @@ function extractGroqAudioInput( span: Span, ): { input: unknown; metadata: Record } { const source = params.file ?? params.url; - const filePart = groqAudioFilePart(source); + const captureAttachments = isAutoCaptureAttachmentsEnabled(span); + const filePart = groqAudioFilePart(source, captureAttachments); const input = { operation, ...(params.prompt !== undefined ? { prompt: params.prompt } : {}), @@ -160,16 +164,37 @@ function extractGroqAudioInput( : {}), }, }; - if (!filePart) observeGroqAudioInput(source, input, span); + if (!filePart && captureAttachments) + observeGroqAudioInput(source, input, span); return { input, metadata: { model: params.model, provider: "groq" }, }; } -function groqAudioFilePart(source: unknown): unknown | undefined { +function groqAudioFilePart( + source: unknown, + captureAttachments: boolean, +): unknown | undefined { if (source === undefined || source === null) return undefined; + if (!captureAttachments) { + return { + type: "file", + file: { + filename: + typeof source === "string" + ? (filenameFromPath(source) ?? "audio") + : isObject(source) + ? (filenameForAudioSource(source) ?? "audio") + : "audio", + ...(typeof source === "string" && /^https?:\/\//i.test(source) + ? { file_data: source } + : {}), + }, + }; + } + if (typeof source === "string") { if (source.startsWith("data:")) { const contentType = source.match(/^data:([^;,]+)/)?.[1]; @@ -378,22 +403,25 @@ function captureGroqSpeechResponse( span: Span, startTime: number, ): boolean { + const captureAttachments = isAutoCaptureAttachmentsEnabled(span); if (!isObject(response)) return false; - const headerContentType = responseHeader(response, "content-type")?.split( - ";", - 1, - )[0]; - const contentType = headerContentType?.startsWith("audio/") - ? headerContentType - : audioContentTypeFromFormat(request.response_format ?? "wav"); + const headerContentType = captureAttachments + ? responseHeader(response, "content-type")?.split(";", 1)[0] + : undefined; + const contentType = !captureAttachments + ? "application/octet-stream" + : headerContentType?.startsWith("audio/") + ? headerContentType + : audioContentTypeFromFormat(request.response_format ?? "wav"); if (!contentType) return false; - const filename = - filenameFromContentDisposition( - responseHeader(response, "content-disposition"), - ) ?? `speech.${getExtensionFromMediaType(contentType)}`; - const chunks: Uint8Array[] = []; + const filename = !captureAttachments + ? "" + : (filenameFromContentDisposition( + responseHeader(response, "content-disposition"), + ) ?? `speech.${getExtensionFromMediaType(contentType)}`); + const chunks: Uint8Array[] | undefined = captureAttachments ? [] : undefined; let finished = false; let started = false; let spanEnded = false; @@ -407,6 +435,7 @@ function captureGroqSpeechResponse( } }; let firstChunk = true; + let byteSize = 0; const onChunk = (chunk: Uint8Array) => { if (finished || chunk.byteLength === 0) return; if (firstChunk) { @@ -417,29 +446,32 @@ function captureGroqSpeechResponse( }); firstChunk = false; } - chunks.push(new Uint8Array(chunk)); + byteSize += chunk.byteLength; + if (chunks) chunks.push(new Uint8Array(chunk)); }; const complete = () => { if (finished) return; finished = true; - const data = concatUint8Arrays(...chunks); - chunks.length = 0; try { span.log({ output: { content: - data.byteLength > 0 + byteSize > 0 && chunks ? [ { type: "file", file: { filename, - byte_size: data.byteLength, - file_data: new Attachment({ - data: data.buffer, - filename, - contentType, - }), + ...(byteSize ? { byte_size: byteSize } : {}), + ...(chunks + ? { + file_data: new Attachment({ + data: concatUint8Arrays(...chunks).buffer, + filename, + contentType, + }), + } + : {}), }, }, ] @@ -447,30 +479,35 @@ function captureGroqSpeechResponse( }, }); } finally { + if (chunks) chunks.length = 0; endSpan(); } }; const cancel = (error?: unknown) => { if (finished) return; finished = true; - chunks.length = 0; try { if (error !== undefined) span.log({ error }); } finally { + if (chunks) chunks.length = 0; endSpan(); } }; - const patched = observeResponseBytes(response, { - onChunk, - onComplete: complete, - onCancel: cancel, - onStart: () => { - started = true; + const patched = observeResponseBytes( + response, + { + onChunk, + onComplete: complete, + onCancel: cancel, + onStart: () => { + started = true; + }, + aroundRead: (next) => withCurrent(span, next), + debugLabel: "Groq speech audio", }, - aroundRead: (next) => withCurrent(span, next), - debugLabel: "Groq speech audio", - }); + captureAttachments, + ); if (patched) queueMicrotask(() => { if (!started) { @@ -489,6 +526,7 @@ function captureGroqSpeechResponse( function observeResponseBytes( response: Record, options: Parameters[1], + captureAttachments = true, ): boolean { const onStart = options.onStart; const safely = (callback: () => void) => { @@ -518,11 +556,15 @@ function observeResponseBytes( patched = patchResponseBytesMethod(response, "arrayBuffer", safeOptions) || patched; patched = patchResponseBytesMethod(response, "bytes", safeOptions) || patched; - patched = patchResponseBlobMethod(response, safeOptions) || patched; + patched = + patchResponseBlobMethod(response, safeOptions, captureAttachments) || + patched; for (const method of ["formData", "json", "text"] as const) patched = patchResponseCompletionMethod(response, method, safeOptions) || patched; - patched = patchResponseCloneMethod(response, safeOptions) || patched; + patched = + patchResponseCloneMethod(response, safeOptions, captureAttachments) || + patched; return patched; } @@ -557,6 +599,7 @@ function patchResponseBytesMethod( function patchResponseBlobMethod( response: object, options: Parameters[1], + captureAttachments: boolean, ): boolean { const responseRecord = response as Record; const original = responseRecord.blob; @@ -572,7 +615,7 @@ function patchResponseBlobMethod( throw error; } void Promise.resolve(result).then((value) => { - if (!(value instanceof Blob)) { + if (!captureAttachments || !(value instanceof Blob)) { options.onComplete(); return; } @@ -616,6 +659,7 @@ function patchResponseCompletionMethod( function patchResponseCloneMethod( response: object, options: Parameters[1], + captureAttachments: boolean, ): boolean { const responseRecord = response as Record; const original = responseRecord.clone; @@ -623,7 +667,8 @@ function patchResponseCloneMethod( return false; responseRecord.clone = function (this: unknown, ...args: unknown[]) { const clone = Reflect.apply(original, this, args); - if (isObject(clone)) observeResponseBytes(clone, options); + if (isObject(clone)) + observeResponseBytes(clone, options, captureAttachments); return clone; }; return true; diff --git a/js/src/instrumentation/plugins/ollama-plugin.test.ts b/js/src/instrumentation/plugins/ollama-plugin.test.ts index 0006efea8..11c5b147f 100644 --- a/js/src/instrumentation/plugins/ollama-plugin.test.ts +++ b/js/src/instrumentation/plugins/ollama-plugin.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-plugin"; 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/plugins/ollama-plugin.ts b/js/src/instrumentation/plugins/ollama-plugin.ts index f48040348..2c411250e 100644 --- a/js/src/instrumentation/plugins/ollama-plugin.ts +++ b/js/src/instrumentation/plugins/ollama-plugin.ts @@ -1,3 +1,4 @@ +import { isAutoCaptureAttachmentsEnabled } from "../../wrappers/attachment-utils"; import { SpanTypeAttribute, isObject } from "../../../util/index"; import iso from "../../isomorph"; import { Attachment } from "../../logger"; @@ -209,7 +210,17 @@ function normalizeTextAndImages( const imageParts: Record[] = []; const unrecognizedImages: unknown[] = []; + const captureAttachments = isAutoCaptureAttachmentsEnabled(); 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/plugins/openai-media.test.ts b/js/src/instrumentation/plugins/openai-media.test.ts index 2232b5945..e12581dc9 100644 --- a/js/src/instrumentation/plugins/openai-media.test.ts +++ b/js/src/instrumentation/plugins/openai-media.test.ts @@ -176,7 +176,7 @@ it("does not read media upload values without attachment opt-in", async () => { expect(reads).toBe(0); expect( rows.find((row) => row.input?.content)?.input?.content?.[0]?.image_url?.url, - ).toBe(""); + ).toBeUndefined(); }); it("copies consumed audio before application mutations", async () => { diff --git a/js/src/instrumentation/plugins/openai-media.ts b/js/src/instrumentation/plugins/openai-media.ts index 3e20aec1f..10d31eacf 100644 --- a/js/src/instrumentation/plugins/openai-media.ts +++ b/js/src/instrumentation/plugins/openai-media.ts @@ -32,13 +32,13 @@ type MediaPart = | { type: "text"; text: string } | { type: "image_url"; - image_url: { url: unknown }; + image_url?: { url: unknown }; purpose?: string; revised_prompt?: string; } | { type: "file"; - file: { filename: string; file_data: unknown; byte_size?: number }; + file: { filename: string; file_data?: unknown; byte_size?: number }; }; const AUDIO_TYPES = new Map( @@ -60,7 +60,7 @@ function mediaAttachment( ): unknown { if (value instanceof URL) return value.toString(); if (typeof value === "string" && /^https?:/.test(value)) return value; - if (!captureAttachments) return ""; + if (!captureAttachments) return undefined; const blob = value instanceof Blob ? value : convertDataToBlob(value, contentType); if (blob) @@ -109,7 +109,7 @@ async function mediaInput( operation: string, captureAttachments: boolean, ) { - const pendingContent: Array> = []; + const pendingContent: Array> = []; for (const [key, purpose] of [ ["image", "reference"], ["mask", "mask"], @@ -118,8 +118,23 @@ async function mediaInput( const value = params[key]; if (value === undefined) continue; for (const item of Array.isArray(value) ? value : [value]) { + if (!captureAttachments && key !== "file") { + if ( + item instanceof URL || + (typeof item === "string" && /^https?:/.test(item)) + ) { + pendingContent.push( + Promise.resolve({ + type: "image_url", + image_url: { url: item instanceof URL ? item.toString() : item }, + purpose, + }), + ); + } + continue; + } pendingContent.push( - (async (): Promise => { + (async (): Promise => { const isImage = key !== "file"; const sourceName = isObject(item) && typeof item.name === "string" @@ -213,9 +228,18 @@ async function mediaInput( filename, captureAttachments, ); + if (isImage && attachment === undefined) return undefined; return isImage ? { type: "image_url", image_url: { url: attachment }, purpose } - : { type: "file", file: { filename, file_data: attachment } }; + : { + type: "file", + file: { + filename, + ...(attachment !== undefined + ? { file_data: attachment } + : {}), + }, + }; })(), ); } @@ -236,7 +260,9 @@ async function mediaInput( ) parameters.format = params.response_format; const prompt = params.prompt ?? params.input; - const content = await Promise.all(pendingContent); + const content = (await Promise.all(pendingContent)).filter( + (part) => part !== undefined, + ); return { operation, ...(prompt !== undefined ? { prompt } : {}), @@ -263,10 +289,10 @@ function mediaOutput( captureAttachments, ) : item.url; - if (url !== undefined) + if (url !== undefined || item.revised_prompt) content.push({ type: "image_url", - image_url: { url }, + ...(url !== undefined ? { image_url: { url } } : {}), ...(item.revised_prompt ? { revised_prompt: item.revised_prompt } : {}), }); } @@ -753,16 +779,20 @@ export function interceptOpenAIMedia( const accumulated: OpenAIMediaResult = {}; const audio: Blob[] = []; patchStreamIfNeeded(value, { - onChunk: (event) => { + shouldCollect: (event) => { if (event.b64_json || event.audio || event.delta || event.text) first(); if (event.usage) accumulated.usage = event.usage; if (event.model) accumulated.model = event.model; - if (event.type.endsWith(".completed") && event.b64_json) { + if ( + captureAttachments && + event.type.endsWith(".completed") && + event.b64_json + ) { accumulated.data = [ ...(accumulated.data ?? []), { - b64_json: captureAttachments ? event.b64_json : "", + b64_json: event.b64_json, }, ]; accumulated.output_format = event.output_format; @@ -782,6 +812,7 @@ export function interceptOpenAIMedia( ); if (blob) audio.push(blob); } + return false; }, onComplete: () => { finish(accumulated); diff --git a/js/src/instrumentation/plugins/openai-plugin.test.ts b/js/src/instrumentation/plugins/openai-plugin.test.ts index ddcb66e7d..5cde40aa4 100644 --- a/js/src/instrumentation/plugins/openai-plugin.test.ts +++ b/js/src/instrumentation/plugins/openai-plugin.test.ts @@ -1213,9 +1213,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", () => { diff --git a/js/src/instrumentation/plugins/openai-span-data.ts b/js/src/instrumentation/plugins/openai-span-data.ts index a1414ca84..5ff0bf913 100644 --- a/js/src/instrumentation/plugins/openai-span-data.ts +++ b/js/src/instrumentation/plugins/openai-span-data.ts @@ -2,6 +2,7 @@ import { Attachment } from "../../logger"; import { isAutoCaptureAttachmentsEnabled, + omitMediaData, processInputAttachments, } from "../../wrappers/attachment-utils"; import { isObject } from "../../../util/index"; @@ -101,9 +102,9 @@ export function processImagesInOutput( captureAttachments = isAutoCaptureAttachmentsEnabled(), ): 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 +114,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/plugins/pi-coding-agent-plugin.test.ts b/js/src/instrumentation/plugins/pi-coding-agent-plugin.test.ts index 73b5afd24..bbff5c587 100644 --- a/js/src/instrumentation/plugins/pi-coding-agent-plugin.test.ts +++ b/js/src/instrumentation/plugins/pi-coding-agent-plugin.test.ts @@ -26,6 +26,9 @@ vi.mock("../../isomorph", async (importOriginal) => { }); vi.mock("../../logger", () => ({ + CAPTURE_ATTACHMENTS: Symbol.for("braintrust.captureAttachments"), + _internalGetGlobalState: () => undefined, + BaseAttachment: class {}, startSpan: (...args: unknown[]) => mockStartSpan(...args), withCurrent: (_span: unknown, callback: () => unknown) => callback(), })); diff --git a/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts b/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts index 174ded6ce..b8aad9abf 100644 --- a/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts +++ b/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts @@ -2,7 +2,11 @@ import { BasePlugin, toLoggedError } from "../core"; import type { ChannelMessage } from "../core/channel-definitions"; import iso, { type IsoAsyncLocalStorage } from "../../isomorph"; import { debugLogger } from "../../debug-logger"; -import { startSpan as startBaseSpan, withCurrent } from "../../logger"; +import { + CAPTURE_ATTACHMENTS, + startSpan as startBaseSpan, + withCurrent, +} from "../../logger"; import type { Span } from "../../logger"; import { INSTRUMENTATION_NAMES, @@ -10,7 +14,10 @@ import { } from "../../span-origin"; import { getCurrentUnixTimestamp } from "../../util"; import { SpanTypeAttribute, isObject } from "../../../util/index"; -import { processInputAttachments } from "../../wrappers/attachment-utils"; +import { + isAutoCaptureAttachmentsEnabled, + processInputAttachments, +} from "../../wrappers/attachment-utils"; import { runWithAutoInstrumentationAllowed, runWithAutoInstrumentationSuppressed, @@ -392,10 +399,14 @@ async function startPiLlmSpan( withSpanInstrumentationName( { event: { - input: processInputAttachments(normalizePiContextInput(context)), + input: processInputAttachments( + normalizePiContextInput(context), + isAutoCaptureAttachmentsEnabled(state.span), + ), metadata, }, name: getLlmSpanName(model), + [CAPTURE_ATTACHMENTS]: isAutoCaptureAttachmentsEnabled(state.span), parent: await state.span.export(), spanAttributes: { type: SpanTypeAttribute.LLM }, }, @@ -560,13 +571,19 @@ async function handlePiAgentEvent( switch (event.type) { case "message_end": if (isPiAssistantMessage(event.message)) { - state.output = extractAssistantOutput(event.message); + state.output = extractAssistantOutput( + event.message, + isAutoCaptureAttachmentsEnabled(state.span), + ); } return; case "turn_end": state.turnEnded = true; if (isPiAssistantMessage(event.message)) { - state.output = extractAssistantOutput(event.message); + state.output = extractAssistantOutput( + event.message, + isAutoCaptureAttachmentsEnabled(state.span), + ); if (!state.collectedLlmUsageMetrics) { addMetrics(state.metrics, extractUsageMetrics(event.message.usage)); } @@ -607,6 +624,7 @@ async function startPiToolSpan( metadata, }, name: event.toolName || "tool", + [CAPTURE_ATTACHMENTS]: isAutoCaptureAttachmentsEnabled(state.span), parent: await state.span.export(), spanAttributes: { type: SpanTypeAttribute.TOOL }, }, @@ -715,7 +733,13 @@ function finishPiLlmSpan( ...(message ? extractAssistantMetadata(message) : {}), }, metrics, - ...(message ? { output: extractAssistantOutput(message) } : {}), + ...(message + ? { + output: withCurrent(llmState.span, () => + extractAssistantOutput(message), + ), + } + : {}), }); } finally { llmState.span.end(); @@ -823,14 +847,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 = isAutoCaptureAttachmentsEnabled(), +): unknown { + return processInputAttachments( + [ + { + finish_reason: normalizeStopReason(message.stopReason), + index: 0, + message: normalizeAssistantMessage(message), + }, + ], + captureAttachments, + ); } function isPiUserMessage( diff --git a/js/src/instrumentation/plugins/strands-agent-sdk-plugin.test.ts b/js/src/instrumentation/plugins/strands-agent-sdk-plugin.test.ts index a2b3f0527..e7efb61c5 100644 --- a/js/src/instrumentation/plugins/strands-agent-sdk-plugin.test.ts +++ b/js/src/instrumentation/plugins/strands-agent-sdk-plugin.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"; @@ -321,11 +326,18 @@ describe("StrandsAgentSDKPlugin", () => { }); 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, + ); const plugin = new StrandsAgentSDKPlugin(); plugin.enable(); @@ -372,10 +384,16 @@ describe("StrandsAgentSDKPlugin", () => { (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/plugins/strands-agent-sdk-plugin.ts b/js/src/instrumentation/plugins/strands-agent-sdk-plugin.ts index 13a434cff..b27e7cbe1 100644 --- a/js/src/instrumentation/plugins/strands-agent-sdk-plugin.ts +++ b/js/src/instrumentation/plugins/strands-agent-sdk-plugin.ts @@ -15,7 +15,11 @@ import { import { LRUCache } from "../../lru-cache"; import { getCurrentUnixTimestamp } from "../../util"; import { SpanTypeAttribute, isObject } from "../../../util/index"; -import { convertDataToBlob } from "../../wrappers/attachment-utils"; +import { + convertDataToBlob, + isAutoCaptureAttachmentsEnabled, + 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>; + }; export class StrandsAgentSDKPlugin extends BasePlugin { private readonly activeChildParents: ActiveChildParents = new WeakMap(); @@ -246,7 +253,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, () => @@ -788,7 +795,8 @@ function finalizeAgentStream( ...(output !== undefined ? { output } : {}), }); state.span.end(); - state.attachmentCache.strings.clear(); + if (state.attachmentCache.captureAttachments) + state.attachmentCache.strings.clear(); } function finalizeMultiAgentStream( @@ -1013,8 +1021,11 @@ const STRANDS_MEDIA_TYPES: Record = { xml: "application/xml", }; -function createStrandsAttachmentCache(): StrandsAttachmentCache { +function createStrandsAttachmentCache(parent?: Span): StrandsAttachmentCache { + if (!isAutoCaptureAttachmentsEnabled(parent)) + return { captureAttachments: false }; return { + captureAttachments: true, objects: new WeakMap(), strings: new LRUCache({ max: MAX_STRANDS_STRING_ATTACHMENT_CACHE_ENTRIES, @@ -1141,35 +1152,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 } : {}), + }, + } + : {}), }, }; } @@ -1180,6 +1200,7 @@ function getOrCreateStrandsAttachment( contentType: string, cache: StrandsAttachmentCache, ): Attachment | undefined { + if (!cache.captureAttachments) return undefined; const key = `${contentType}\0${filename}`; const attachments = typeof data === "string" diff --git a/js/src/instrumentation/plugins/voyageai-plugin.test.ts b/js/src/instrumentation/plugins/voyageai-plugin.test.ts index 65cfca93c..ff225344a 100644 --- a/js/src/instrumentation/plugins/voyageai-plugin.test.ts +++ b/js/src/instrumentation/plugins/voyageai-plugin.test.ts @@ -21,6 +21,7 @@ describe("VoyageAIPlugin", () => { beforeEach(() => { backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); initLogger({ + captureAttachments: true, projectName: "voyageai-plugin.test.ts", projectId: "test-project-id", }); diff --git a/js/src/logger.ts b/js/src/logger.ts index 5ae72b2fc..afa3516d6 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -715,6 +715,9 @@ function normalizeProxyConnUrl(proxyUrl: string): string { : proxyUrl; } +/** @internal */ +export const CAPTURE_ATTACHMENTS = Symbol.for("braintrust.captureAttachments"); + export class BraintrustState { public id: string; public currentExperiment: Experiment | undefined; @@ -739,6 +742,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; @@ -1169,6 +1174,19 @@ export class BraintrustState { this.bgLogger().setMaskingFunction(maskingFunction); } + /** @internal Resolve capture against the actual logger or span, including across SDK bundles. */ + public _internalCaptureAttachmentsEnabled(parent?: object): boolean { + const policy = parent && Reflect.get(parent, CAPTURE_ATTACHMENTS); + if (typeof policy === "boolean") return policy; + return ( + this.captureAttachments ?? + ["1", "true"].includes( + iso.getEnv("BRAINTRUST_CAPTURE_ATTACHMENTS")?.trim().toLowerCase() ?? + "", + ) + ); + } + public setDebugLogLevel(option: DebugLogLevelOption): void { if (option === undefined) { return; @@ -2751,6 +2769,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; @@ -2758,7 +2777,9 @@ export class Logger implements Exportable { constructor( state: BraintrustState, lazyMetadata: LazyValue, - logOptions: LogOptions = {}, + logOptions: LogOptions & { + captureAttachments?: boolean; + } = {}, ) { this.lazyMetadata = lazyMetadata; this._asyncFlush = logOptions.asyncFlush; @@ -2767,6 +2788,7 @@ export class Logger implements Exportable { this.lastStartTime = getCurrentUnixTimestamp(); this.lazyId = new LazyValue(async () => await this.id); this.calledStartSpan = false; + this.captureAttachments = logOptions.captureAttachments; this.state = state; } @@ -2790,6 +2812,17 @@ export class Logger implements Exportable { return this.state; } + /** @internal */ + public get [CAPTURE_ATTACHMENTS](): boolean { + return ( + this.captureAttachments ?? + ["1", "true"].includes( + iso.getEnv("BRAINTRUST_CAPTURE_ATTACHMENTS")?.trim().toLowerCase() ?? + "", + ) + ); + } + private parentObjectType() { return SpanObjectTypeV3.PROJECT_LOGS; } @@ -2887,6 +2920,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. @@ -4975,6 +5009,15 @@ type AsyncFlushArg = { }; 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; environment?: SpanOriginEnvironment; @@ -5053,7 +5096,16 @@ export function initLogger( }, ); + if ( + (options.setCurrent ?? true) && + options.captureAttachments !== undefined + ) { + state.captureAttachments = options.captureAttachments; + } const ret = new Logger(state, lazyMetadata, { + captureAttachments: + options.captureAttachments ?? + ((options.setCurrent ?? true) ? state.captureAttachments : undefined), asyncFlush, computeMetadataArgs, linkArgs, @@ -8107,6 +8159,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; @@ -8143,11 +8197,15 @@ 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._internalCaptureAttachmentsEnabled(); this._propagatedState = args.propagatedState; const instrumentationName = getSpanInstrumentationName(args) ?? @@ -8381,6 +8439,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, @@ -8406,6 +8465,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/wrappers/attachment-utils.test.ts b/js/src/wrappers/attachment-utils.test.ts index 3994ecb5c..f8edface1 100644 --- a/js/src/wrappers/attachment-utils.test.ts +++ b/js/src/wrappers/attachment-utils.test.ts @@ -1,4 +1,4 @@ -import { afterEach, describe, expect, it } from "vitest"; +import { afterEach, describe, expect, it, vi } from "vitest"; import iso from "../isomorph"; import { isAutoCaptureAttachmentsEnabled, @@ -50,7 +50,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 +58,32 @@ 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); +}); diff --git a/js/src/wrappers/attachment-utils.ts b/js/src/wrappers/attachment-utils.ts index 47876ddf7..1a51108d3 100644 --- a/js/src/wrappers/attachment-utils.ts +++ b/js/src/wrappers/attachment-utils.ts @@ -1,13 +1,41 @@ -import { Attachment } from "../logger"; import iso from "../isomorph"; +import { + Attachment, + BaseAttachment, + _internalGetGlobalState, + getSpanParentObject, +} from "../logger"; -const CAPTURE_ATTACHMENTS_ENV_VAR = "BRAINTRUST_CAPTURE_ATTACHMENTS"; +export function isAutoCaptureAttachmentsEnabled(parent?: object): boolean { + const state = _internalGetGlobalState(); + return state + ? state._internalCaptureAttachmentsEnabled(parent ?? getSpanParentObject()) + : ["1", "true"].includes( + iso.getEnv("BRAINTRUST_CAPTURE_ATTACHMENTS")?.trim().toLowerCase() ?? + "", + ); +} -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; } /** @@ -90,7 +118,7 @@ export function convertDataToBlob(data: any, mediaType: string): Blob | null { */ export function processInputAttachments( input: any, - captureAttachments = true, + captureAttachments = isAutoCaptureAttachmentsEnabled(), ): any { if (!input) { return input; @@ -125,13 +153,15 @@ 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 URL) return node; + // OpenAI chat image_url content format if ( node.type === "image_url" && @@ -140,14 +170,17 @@ export function processInputAttachments( typeof node.image_url.url === "string" && node.image_url.url.startsWith("data:") ) { + if (!captureAttachments) + return omitMediaData({ + ...node, + image_url: omitMediaData(node.image_url, "url"), + }); const mediaType = inferMediaTypeFromDataUrl( node.image_url.url, "image/png", ); const filename = `image.${getExtensionFromMediaType(mediaType)}`; - const attachment = captureAttachments - ? toAttachment(node.image_url.url, mediaType, filename) - : ""; + const attachment = toAttachment(node.image_url.url, mediaType, filename); if (attachment) { return { @@ -179,14 +212,13 @@ 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", ); const filename = `${node.type === "video_base64" ? "video" : "image"}.${getExtensionFromMediaType(mediaType)}`; - const attachment = captureAttachments - ? toAttachment(voyageBase64Value, mediaType, filename) - : ""; + const attachment = toAttachment(voyageBase64Value, mediaType, filename); if (attachment) { return { @@ -204,6 +236,11 @@ export function processInputAttachments( typeof node.file.file_data === "string" && node.file.file_data.startsWith("data:") ) { + if (!captureAttachments) + return omitMediaData({ + ...node, + file: omitMediaData(node.file, "file_data"), + }); const mediaType = inferMediaTypeFromDataUrl( node.file.file_data, "application/octet-stream", @@ -212,9 +249,7 @@ export function processInputAttachments( typeof node.file.filename === "string" && node.file.filename ? node.file.filename : `document.${getExtensionFromMediaType(mediaType)}`; - const attachment = captureAttachments - ? toAttachment(node.file.file_data, mediaType, filename) - : ""; + const attachment = toAttachment(node.file.file_data, mediaType, filename); if (attachment) { return { @@ -229,6 +264,16 @@ export function processInputAttachments( // AI SDK image content format if (node.type === "image" && node.image) { + if (node.image instanceof BaseAttachment) { + if (captureAttachments) attachmentIndex++; + return node; + } + if (!captureAttachments) { + return node.image instanceof URL || + (typeof node.image === "string" && /^https?:/.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); @@ -237,12 +282,7 @@ export function processInputAttachments( } const filename = `input_image_${attachmentIndex}.${getExtensionFromMediaType(mediaType)}`; - const attachment = captureAttachments - ? toAttachment(node.image, mediaType, filename) - : node.image instanceof URL || - (typeof node.image === "string" && /^https?:/.test(node.image)) - ? null - : ""; + const attachment = toAttachment(node.image, mediaType, filename); if (attachment) { attachmentIndex++; @@ -255,16 +295,21 @@ export function processInputAttachments( // AI SDK file content format if (node.type === "file" && node.data) { + if (node.data instanceof BaseAttachment) { + if (captureAttachments) attachmentIndex++; + return node; + } + if (!captureAttachments) { + return node.data instanceof URL || + (typeof node.data === "string" && /^https?:/.test(node.data)) + ? node + : omitMediaData(node, "data"); + } const mediaType = node.mediaType || "application/octet-stream"; const filename = node.filename || `input_file_${attachmentIndex}.${getExtensionFromMediaType(mediaType)}`; - const attachment = captureAttachments - ? toAttachment(node.data, mediaType, filename) - : node.data instanceof URL || - (typeof node.data === "string" && /^https?:/.test(node.data)) - ? null - : ""; + const attachment = toAttachment(node.data, mediaType, filename); if (attachment) { attachmentIndex++; @@ -277,13 +322,14 @@ export function processInputAttachments( 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..1e5c5c7f9 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", }); diff --git a/js/tests/provider-wrappers.test.ts b/js/tests/provider-wrappers.test.ts index 37161415c..26842e07c 100644 --- a/js/tests/provider-wrappers.test.ts +++ b/js/tests/provider-wrappers.test.ts @@ -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) {