From 2c5d66bf347abbc30e01fbdf6f2a28aeab1c0aa8 Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Mon, 20 Apr 2026 20:10:09 +0200 Subject: [PATCH 1/4] feat: Add instrumentation for `groq-sdk` --- .env.example | 1 + .github/workflows/e2e-canary.yaml | 1 + .github/workflows/integration-tests.yaml | 2 + e2e/README.md | 1 + e2e/helpers/scenario-installer.ts | 1 + .../groq-v1-auto.span-events.json | 138 +++++++++++++ .../groq-v1-wrapped.span-events.json | 138 +++++++++++++ .../groq-instrumentation/assertions.ts | 166 ++++++++++++++++ .../groq-instrumentation/constants.mjs | 3 + .../groq-instrumentation/package.json | 14 ++ .../groq-instrumentation/pnpm-lock.yaml | 23 +++ .../groq-instrumentation/scenario.impl.mjs | 106 ++++++++++ .../groq-instrumentation/scenario.mjs | 9 + .../groq-instrumentation/scenario.test.ts | 49 +++++ .../groq-instrumentation/scenario.ts | 9 + e2e/scripts/run-canary-tests-docker.mjs | 1 + .../auto-instrumentations/bundler/plugin.ts | 2 + .../bundler/webpack-loader.ts | 2 + js/src/auto-instrumentations/configs/groq.ts | 31 +++ js/src/auto-instrumentations/hook.mts | 2 + js/src/auto-instrumentations/index.ts | 1 + js/src/exports.ts | 1 + .../instrumentation/braintrust-plugin.test.ts | 44 ++++ js/src/instrumentation/braintrust-plugin.ts | 13 ++ .../instrumentation/plugins/groq-channels.ts | 33 +++ .../plugins/groq-plugin.test.ts | 34 ++++ js/src/instrumentation/plugins/groq-plugin.ts | 127 ++++++++++++ js/src/vendor-sdk-types/groq.ts | 104 ++++++++++ js/src/wrappers/groq.test.ts | 188 ++++++++++++++++++ js/src/wrappers/groq.ts | 159 +++++++++++++++ turbo.json | 8 + 31 files changed, 1411 insertions(+) create mode 100644 e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-auto.span-events.json create mode 100644 e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-wrapped.span-events.json create mode 100644 e2e/scenarios/groq-instrumentation/assertions.ts create mode 100644 e2e/scenarios/groq-instrumentation/constants.mjs create mode 100644 e2e/scenarios/groq-instrumentation/package.json create mode 100644 e2e/scenarios/groq-instrumentation/pnpm-lock.yaml create mode 100644 e2e/scenarios/groq-instrumentation/scenario.impl.mjs create mode 100644 e2e/scenarios/groq-instrumentation/scenario.mjs create mode 100644 e2e/scenarios/groq-instrumentation/scenario.test.ts create mode 100644 e2e/scenarios/groq-instrumentation/scenario.ts create mode 100644 js/src/auto-instrumentations/configs/groq.ts create mode 100644 js/src/instrumentation/plugins/groq-channels.ts create mode 100644 js/src/instrumentation/plugins/groq-plugin.test.ts create mode 100644 js/src/instrumentation/plugins/groq-plugin.ts create mode 100644 js/src/vendor-sdk-types/groq.ts create mode 100644 js/src/wrappers/groq.test.ts create mode 100644 js/src/wrappers/groq.ts diff --git a/.env.example b/.env.example index a44259148..c9c314bb3 100644 --- a/.env.example +++ b/.env.example @@ -6,3 +6,4 @@ OPENROUTER_API_KEY= MISTRAL_API_KEY= HUGGINGFACE_API_KEY= COHERE_API_KEY= +GROQ_API_KEY= diff --git a/.github/workflows/e2e-canary.yaml b/.github/workflows/e2e-canary.yaml index 060e629df..12a287004 100644 --- a/.github/workflows/e2e-canary.yaml +++ b/.github/workflows/e2e-canary.yaml @@ -35,6 +35,7 @@ jobs: BRAINTRUST_E2E_PROJECT_NAME: ${{ vars.BRAINTRUST_E2E_PROJECT_NAME }} GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }} COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }} + GROQ_API_KEY: ${{ secrets.GROQ_API_KEY }} OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} OPENROUTER_API_KEY: ${{ secrets.OPENROUTER_API_KEY }} MISTRAL_API_KEY: ${{ secrets.MISTRAL_API_KEY }} diff --git a/.github/workflows/integration-tests.yaml b/.github/workflows/integration-tests.yaml index f8646475f..7b7b88ef3 100644 --- a/.github/workflows/integration-tests.yaml +++ b/.github/workflows/integration-tests.yaml @@ -58,6 +58,7 @@ jobs: ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }} GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }} COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }} + GROQ_API_KEY: ${{ secrets.GROQ_API_KEY }} OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} OPENROUTER_API_KEY: ${{ secrets.OPENROUTER_API_KEY }} MISTRAL_API_KEY: ${{ secrets.MISTRAL_API_KEY }} @@ -109,6 +110,7 @@ jobs: ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }} GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }} COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }} + GROQ_API_KEY: ${{ secrets.GROQ_API_KEY }} OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} OPENROUTER_API_KEY: ${{ secrets.OPENROUTER_API_KEY }} MISTRAL_API_KEY: ${{ secrets.MISTRAL_API_KEY }} diff --git a/e2e/README.md b/e2e/README.md index b165f3258..b4a15d56e 100644 --- a/e2e/README.md +++ b/e2e/README.md @@ -119,6 +119,7 @@ Non-hermetic scenarios require provider credentials in addition to the mock Brai - `MISTRAL_API_KEY` - `HUGGINGFACE_API_KEY` - `COHERE_API_KEY` +- `GROQ_API_KEY` `claude-agent-sdk-instrumentation` also uses `ANTHROPIC_API_KEY`, because it runs the real Claude Agent SDK against Anthropic in the same style as the existing live Anthropic wrapper coverage. diff --git a/e2e/helpers/scenario-installer.ts b/e2e/helpers/scenario-installer.ts index 13062e02b..29318aeff 100644 --- a/e2e/helpers/scenario-installer.ts +++ b/e2e/helpers/scenario-installer.ts @@ -25,6 +25,7 @@ const INSTALL_SECRET_ENV_VARS = [ "GEMINI_API_KEY", "GITHUB_TOKEN", "GH_TOKEN", + "GROQ_API_KEY", "HUGGINGFACE_API_KEY", "OPENAI_API_KEY", "OPENROUTER_API_KEY", diff --git a/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-auto.span-events.json b/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-auto.span-events.json new file mode 100644 index 000000000..edc29e4f9 --- /dev/null +++ b/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-auto.span-events.json @@ -0,0 +1,138 @@ +[ + { + "has_input": false, + "has_output": false, + "metadata": { + "scenario": "groq-instrumentation" + }, + "metric_keys": [], + "name": "groq-instrumentation-root", + "root_span_id": "", + "span_id": "", + "span_parents": [], + "type": "task" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "chat" + }, + "metric_keys": [], + "name": "groq-chat-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "metric_keys": [ + "completion_time", + "completion_tokens", + "prompt_time", + "prompt_tokens", + "queue_time", + "time_to_first_token", + "tokens", + "total_time" + ], + "name": "groq.chat.completions.create", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "stream" + }, + "metric_keys": [], + "name": "groq-stream-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "metric_keys": [ + "completion_time", + "completion_tokens", + "prompt_time", + "prompt_tokens", + "queue_time", + "time_to_first_token", + "tokens", + "total_time" + ], + "name": "groq.chat.completions.create", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "tool" + }, + "metric_keys": [], + "name": "groq-tool-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "metric_keys": [ + "completion_time", + "completion_tokens", + "prompt_time", + "prompt_tokens", + "queue_time", + "time_to_first_token", + "tokens", + "total_time" + ], + "name": "groq.chat.completions.create", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + } +] diff --git a/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-wrapped.span-events.json b/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-wrapped.span-events.json new file mode 100644 index 000000000..edc29e4f9 --- /dev/null +++ b/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-wrapped.span-events.json @@ -0,0 +1,138 @@ +[ + { + "has_input": false, + "has_output": false, + "metadata": { + "scenario": "groq-instrumentation" + }, + "metric_keys": [], + "name": "groq-instrumentation-root", + "root_span_id": "", + "span_id": "", + "span_parents": [], + "type": "task" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "chat" + }, + "metric_keys": [], + "name": "groq-chat-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "metric_keys": [ + "completion_time", + "completion_tokens", + "prompt_time", + "prompt_tokens", + "queue_time", + "time_to_first_token", + "tokens", + "total_time" + ], + "name": "groq.chat.completions.create", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "stream" + }, + "metric_keys": [], + "name": "groq-stream-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "metric_keys": [ + "completion_time", + "completion_tokens", + "prompt_time", + "prompt_tokens", + "queue_time", + "time_to_first_token", + "tokens", + "total_time" + ], + "name": "groq.chat.completions.create", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "tool" + }, + "metric_keys": [], + "name": "groq-tool-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "metric_keys": [ + "completion_time", + "completion_tokens", + "prompt_time", + "prompt_tokens", + "queue_time", + "time_to_first_token", + "tokens", + "total_time" + ], + "name": "groq.chat.completions.create", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + } +] diff --git a/e2e/scenarios/groq-instrumentation/assertions.ts b/e2e/scenarios/groq-instrumentation/assertions.ts new file mode 100644 index 000000000..d3472114b --- /dev/null +++ b/e2e/scenarios/groq-instrumentation/assertions.ts @@ -0,0 +1,166 @@ +import { beforeAll, describe, expect, test } from "vitest"; +import type { Json } from "../../helpers/normalize"; +import type { CapturedLogEvent } from "../../helpers/mock-braintrust-server"; +import { + formatJsonFileSnapshot, + resolveFileSnapshotPath, +} from "../../helpers/file-snapshot"; +import { withScenarioHarness } from "../../helpers/scenario-harness"; +import { findChildSpans, findLatestSpan } from "../../helpers/trace-selectors"; +import { summarizeWrapperContract } from "../../helpers/wrapper-contract"; +import { ROOT_NAME, SCENARIO_NAME } from "./constants.mjs"; + +type RunGroqScenario = (harness: { + runNodeScenarioDir: (options: { + entry: string; + nodeArgs: string[]; + runContext?: { variantKey: string }; + scenarioDir: string; + timeoutMs: number; + }) => Promise; + runScenarioDir: (options: { + entry: string; + runContext?: { variantKey: string }; + scenarioDir: string; + timeoutMs: number; + }) => Promise; +}) => Promise; + +function findGroqSpan( + events: CapturedLogEvent[], + parentId: string | undefined, + spanName: string, +) { + const spans = findChildSpans(events, spanName, parentId); + return spans.find((candidate) => candidate.output !== undefined) ?? spans[0]; +} + +function buildSpanSummary(events: CapturedLogEvent[]): Json { + const chatOperation = findLatestSpan(events, "groq-chat-operation"); + const streamOperation = findLatestSpan(events, "groq-stream-operation"); + const toolOperation = findLatestSpan(events, "groq-tool-operation"); + + return [ + findLatestSpan(events, ROOT_NAME), + chatOperation, + findGroqSpan( + events, + chatOperation?.span.id, + "groq.chat.completions.create", + ), + streamOperation, + findGroqSpan( + events, + streamOperation?.span.id, + "groq.chat.completions.create", + ), + toolOperation, + findGroqSpan( + events, + toolOperation?.span.id, + "groq.chat.completions.create", + ), + ].map((event) => + summarizeWrapperContract(event!, [ + "model", + "operation", + "provider", + "scenario", + "temperature", + ]), + ) as Json; +} + +export function defineGroqInstrumentationAssertions(options: { + name: string; + runScenario: RunGroqScenario; + snapshotName: string; + testFileUrl: string; + timeoutMs: number; +}): void { + const spanSnapshotPath = resolveFileSnapshotPath( + options.testFileUrl, + `${options.snapshotName}.span-events.json`, + ); + const testConfig = { + timeout: options.timeoutMs, + }; + + describe(options.name, () => { + let events: CapturedLogEvent[] = []; + + beforeAll(async () => { + await withScenarioHarness(async (harness) => { + await options.runScenario(harness); + events = harness.events(); + }); + }, options.timeoutMs); + + test("captures the scenario root span", testConfig, () => { + const root = findLatestSpan(events, ROOT_NAME); + expect(root).toBeDefined(); + expect(root?.row.metadata).toMatchObject({ + scenario: SCENARIO_NAME, + }); + }); + + test("captures chat and stream spans", testConfig, () => { + const chatOperation = findLatestSpan(events, "groq-chat-operation"); + const chatSpan = findGroqSpan( + events, + chatOperation?.span.id, + "groq.chat.completions.create", + ); + const streamOperation = findLatestSpan(events, "groq-stream-operation"); + const streamSpan = findGroqSpan( + events, + streamOperation?.span.id, + "groq.chat.completions.create", + ); + + expect(chatSpan?.row.metadata).toMatchObject({ + provider: "groq", + }); + expect(chatSpan?.row.metadata?.model).toBeDefined(); + expect(chatSpan?.output).toBeDefined(); + + expect(streamSpan?.row.metadata).toMatchObject({ + provider: "groq", + }); + expect(streamSpan?.row.metadata?.model).toBeDefined(); + expect(streamSpan?.output).toBeDefined(); + expect(streamSpan?.metrics).toMatchObject({ + time_to_first_token: expect.any(Number), + }); + }); + + test("captures tool calling span", testConfig, () => { + const operation = findLatestSpan(events, "groq-tool-operation"); + const span = findGroqSpan( + events, + operation?.span.id, + "groq.chat.completions.create", + ); + + expect(span?.row.metadata).toMatchObject({ + provider: "groq", + }); + expect(span?.row.metadata?.model).toBeDefined(); + expect(span?.output?.[0]?.message?.tool_calls).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + function: expect.objectContaining({ + name: "get_weather", + }), + }), + ]), + ); + }); + + test("matches span snapshot", testConfig, async () => { + await expect( + formatJsonFileSnapshot(buildSpanSummary(events)), + ).toMatchFileSnapshot(spanSnapshotPath); + }); + }); +} diff --git a/e2e/scenarios/groq-instrumentation/constants.mjs b/e2e/scenarios/groq-instrumentation/constants.mjs new file mode 100644 index 000000000..967d98068 --- /dev/null +++ b/e2e/scenarios/groq-instrumentation/constants.mjs @@ -0,0 +1,3 @@ +export const CHAT_MODEL = "llama-3.3-70b-versatile"; +export const ROOT_NAME = "groq-instrumentation-root"; +export const SCENARIO_NAME = "groq-instrumentation"; diff --git a/e2e/scenarios/groq-instrumentation/package.json b/e2e/scenarios/groq-instrumentation/package.json new file mode 100644 index 000000000..1f410a4c4 --- /dev/null +++ b/e2e/scenarios/groq-instrumentation/package.json @@ -0,0 +1,14 @@ +{ + "name": "@braintrust/e2e-groq-instrumentation", + "private": true, + "braintrustScenario": { + "canary": { + "dependencies": { + "groq-sdk": "latest" + } + } + }, + "dependencies": { + "groq-sdk": "1.1.2" + } +} diff --git a/e2e/scenarios/groq-instrumentation/pnpm-lock.yaml b/e2e/scenarios/groq-instrumentation/pnpm-lock.yaml new file mode 100644 index 000000000..d81367d29 --- /dev/null +++ b/e2e/scenarios/groq-instrumentation/pnpm-lock.yaml @@ -0,0 +1,23 @@ +lockfileVersion: '9.0' + +settings: + autoInstallPeers: true + excludeLinksFromLockfile: false + +importers: + + .: + dependencies: + groq-sdk: + specifier: 1.1.2 + version: 1.1.2 + +packages: + + groq-sdk@1.1.2: + resolution: {integrity: sha512-CZO0XUQQDhn43ri1+lZHxZKpb+bGutgTvFmCJtooexiitGmPqhm1hntOT3nCoaq07e+OpeokVnfUs0i/oQuUaQ==} + hasBin: true + +snapshots: + + groq-sdk@1.1.2: {} diff --git a/e2e/scenarios/groq-instrumentation/scenario.impl.mjs b/e2e/scenarios/groq-instrumentation/scenario.impl.mjs new file mode 100644 index 000000000..5c6c79f27 --- /dev/null +++ b/e2e/scenarios/groq-instrumentation/scenario.impl.mjs @@ -0,0 +1,106 @@ +import { wrapGroq } from "braintrust"; +import { + collectAsync, + runOperation, + runTracedScenario, +} from "../../helpers/provider-runtime.mjs"; +import { CHAT_MODEL, ROOT_NAME, SCENARIO_NAME } from "./constants.mjs"; + +export const GROQ_SCENARIO_TIMEOUT_MS = 120_000; + +function getApiKey() { + return process.env.GROQ_API_KEY; +} + +function getWeatherToolDefinition() { + return { + type: "function", + function: { + name: "get_weather", + description: "Get the weather for a city.", + parameters: { + type: "object", + properties: { + location: { + type: "string", + description: "City name.", + }, + }, + required: ["location"], + }, + }, + }; +} + +export async function runGroqInstrumentationScenario(options) { + const apiKey = getApiKey(); + if (!apiKey) { + throw new Error("Expected GROQ_API_KEY to be set for e2e"); + } + + const baseClient = new options.Groq({ + apiKey, + }); + const client = options.decorateClient + ? options.decorateClient(baseClient) + : baseClient; + + await runTracedScenario({ + callback: async () => { + await runOperation("groq-chat-operation", "chat", async () => { + await client.chat.completions.create({ + max_completion_tokens: 12, + messages: [{ role: "user", content: "Reply with exactly OK." }], + model: CHAT_MODEL, + temperature: 0, + }); + }); + + await runOperation("groq-stream-operation", "stream", async () => { + const stream = await client.chat.completions.create({ + messages: [{ role: "user", content: "Reply with exactly STREAM." }], + model: CHAT_MODEL, + stream: true, + temperature: 0, + }); + await collectAsync(stream); + }); + + await runOperation("groq-tool-operation", "tool", async () => { + await client.chat.completions.create({ + messages: [ + { + role: "user", + content: "Check the weather in Vienna and use the weather tool.", + }, + ], + model: CHAT_MODEL, + temperature: 0, + tool_choice: { + type: "function", + function: { + name: "get_weather", + }, + }, + tools: [getWeatherToolDefinition()], + }); + }); + }, + metadata: { + scenario: SCENARIO_NAME, + }, + projectNameBase: "e2e-groq-instrumentation", + rootName: ROOT_NAME, + }); +} + +export async function runWrappedGroqInstrumentation(options) { + await runGroqInstrumentationScenario({ + decorateClient: wrapGroq, + ...options, + }); +} + +export async function runAutoGroqInstrumentation(options) { + await runGroqInstrumentationScenario(options); +} diff --git a/e2e/scenarios/groq-instrumentation/scenario.mjs b/e2e/scenarios/groq-instrumentation/scenario.mjs new file mode 100644 index 000000000..bdd181041 --- /dev/null +++ b/e2e/scenarios/groq-instrumentation/scenario.mjs @@ -0,0 +1,9 @@ +import Groq from "groq-sdk"; +import { runMain } from "../../helpers/provider-runtime.mjs"; +import { runAutoGroqInstrumentation } from "./scenario.impl.mjs"; + +runMain(async () => { + await runAutoGroqInstrumentation({ + Groq, + }); +}); diff --git a/e2e/scenarios/groq-instrumentation/scenario.test.ts b/e2e/scenarios/groq-instrumentation/scenario.test.ts new file mode 100644 index 000000000..849d6f5ec --- /dev/null +++ b/e2e/scenarios/groq-instrumentation/scenario.test.ts @@ -0,0 +1,49 @@ +import { describe } from "vitest"; +import { + prepareScenarioDir, + readInstalledPackageVersion, + resolveScenarioDir, +} from "../../helpers/scenario-harness"; +import { defineGroqInstrumentationAssertions } from "./assertions"; +import { GROQ_SCENARIO_TIMEOUT_MS } from "./scenario.impl.mjs"; + +const scenarioDir = await prepareScenarioDir({ + scenarioDir: resolveScenarioDir(import.meta.url), +}); +const groqSdkVersion = await readInstalledPackageVersion( + scenarioDir, + "groq-sdk", +); + +describe(`groq sdk ${groqSdkVersion}`, () => { + defineGroqInstrumentationAssertions({ + name: "wrapped instrumentation", + runScenario: async ({ runScenarioDir }) => { + await runScenarioDir({ + entry: "scenario.ts", + runContext: { variantKey: "groq-v1-wrapped" }, + scenarioDir, + timeoutMs: GROQ_SCENARIO_TIMEOUT_MS, + }); + }, + snapshotName: "groq-v1-wrapped", + testFileUrl: import.meta.url, + timeoutMs: GROQ_SCENARIO_TIMEOUT_MS, + }); + + defineGroqInstrumentationAssertions({ + name: "auto-hook instrumentation", + runScenario: async ({ runNodeScenarioDir }) => { + await runNodeScenarioDir({ + entry: "scenario.mjs", + nodeArgs: ["--import", "braintrust/hook.mjs"], + runContext: { variantKey: "groq-v1-auto" }, + scenarioDir, + timeoutMs: GROQ_SCENARIO_TIMEOUT_MS, + }); + }, + snapshotName: "groq-v1-auto", + testFileUrl: import.meta.url, + timeoutMs: GROQ_SCENARIO_TIMEOUT_MS, + }); +}); diff --git a/e2e/scenarios/groq-instrumentation/scenario.ts b/e2e/scenarios/groq-instrumentation/scenario.ts new file mode 100644 index 000000000..5a960d526 --- /dev/null +++ b/e2e/scenarios/groq-instrumentation/scenario.ts @@ -0,0 +1,9 @@ +import Groq from "groq-sdk"; +import { runMain } from "../../helpers/provider-runtime.mjs"; +import { runWrappedGroqInstrumentation } from "./scenario.impl.mjs"; + +runMain(async () => { + await runWrappedGroqInstrumentation({ + Groq, + }); +}); diff --git a/e2e/scripts/run-canary-tests-docker.mjs b/e2e/scripts/run-canary-tests-docker.mjs index a2d25868d..0635ba044 100644 --- a/e2e/scripts/run-canary-tests-docker.mjs +++ b/e2e/scripts/run-canary-tests-docker.mjs @@ -19,6 +19,7 @@ const ALLOWED_ENV_KEYS = [ "GEMINI_API_KEY", "GOOGLE_API_KEY", "COHERE_API_KEY", + "GROQ_API_KEY", "OPENAI_API_KEY", "OPENAI_BASE_URL", "OPENROUTER_API_KEY", diff --git a/js/src/auto-instrumentations/bundler/plugin.ts b/js/src/auto-instrumentations/bundler/plugin.ts index 508bb8ae3..8d38d7a57 100644 --- a/js/src/auto-instrumentations/bundler/plugin.ts +++ b/js/src/auto-instrumentations/bundler/plugin.ts @@ -30,6 +30,7 @@ import { openRouterAgentConfigs } from "../configs/openrouter-agent"; import { openRouterConfigs } from "../configs/openrouter"; import { mistralConfigs } from "../configs/mistral"; import { cohereConfigs } from "../configs/cohere"; +import { groqConfigs } from "../configs/groq"; export interface BundlerPluginOptions { /** @@ -81,6 +82,7 @@ export const unplugin = createUnplugin((options = {}) => { ...openRouterAgentConfigs, ...mistralConfigs, ...cohereConfigs, + ...groqConfigs, ...(options.instrumentations || []), ]; diff --git a/js/src/auto-instrumentations/bundler/webpack-loader.ts b/js/src/auto-instrumentations/bundler/webpack-loader.ts index fef1cf416..5d6c7e20b 100644 --- a/js/src/auto-instrumentations/bundler/webpack-loader.ts +++ b/js/src/auto-instrumentations/bundler/webpack-loader.ts @@ -39,6 +39,7 @@ import { openRouterAgentConfigs } from "../configs/openrouter-agent"; import { openRouterConfigs } from "../configs/openrouter"; import { mistralConfigs } from "../configs/mistral"; import { cohereConfigs } from "../configs/cohere"; +import { groqConfigs } from "../configs/groq"; import { type BundlerPluginOptions } from "./plugin"; /** @@ -76,6 +77,7 @@ function getMatcher(options: BundlerPluginOptions): InstrumentationMatcher { ...openRouterAgentConfigs, ...mistralConfigs, ...cohereConfigs, + ...groqConfigs, ...(options.instrumentations ?? []), ]; const dcModule = options.browser ? "dc-browser" : undefined; diff --git a/js/src/auto-instrumentations/configs/groq.ts b/js/src/auto-instrumentations/configs/groq.ts new file mode 100644 index 000000000..9528303de --- /dev/null +++ b/js/src/auto-instrumentations/configs/groq.ts @@ -0,0 +1,31 @@ +import type { InstrumentationConfig } from "@apm-js-collab/code-transformer"; +import { groqChannels } from "../../instrumentation/plugins/groq-channels"; + +export const groqConfigs: InstrumentationConfig[] = [ + { + channelName: groqChannels.chatCompletionsCreate.channelName, + module: { + name: "groq-sdk", + versionRange: ">=1.0.0", + filePath: "resources/chat/completions.mjs", + }, + functionQuery: { + className: "Completions", + methodName: "create", + kind: "Async", + }, + }, + { + channelName: groqChannels.embeddingsCreate.channelName, + module: { + name: "groq-sdk", + versionRange: ">=1.0.0", + filePath: "resources/embeddings.mjs", + }, + functionQuery: { + className: "Embeddings", + methodName: "create", + kind: "Async", + }, + }, +]; diff --git a/js/src/auto-instrumentations/hook.mts b/js/src/auto-instrumentations/hook.mts index 20d6f5a04..2c3622600 100644 --- a/js/src/auto-instrumentations/hook.mts +++ b/js/src/auto-instrumentations/hook.mts @@ -25,6 +25,7 @@ import { openRouterConfigs } from "./configs/openrouter.js"; import { mistralConfigs } from "./configs/mistral.js"; import { googleADKConfigs } from "./configs/google-adk.js"; import { cohereConfigs } from "./configs/cohere.js"; +import { groqConfigs } from "./configs/groq.js"; import { ModulePatch } from "./loader/cjs-patch.js"; import { patchTracingChannel } from "./patch-tracing-channel.js"; @@ -82,6 +83,7 @@ const allConfigs = [ ? [] : googleADKConfigs), ...(isDisabled(disabledIntegrations, "cohere") ? [] : cohereConfigs), + ...(isDisabled(disabledIntegrations, "groq", "groq-sdk") ? [] : groqConfigs), ]; // 1. Register ESM loader for ESM modules diff --git a/js/src/auto-instrumentations/index.ts b/js/src/auto-instrumentations/index.ts index b73a24a02..03fda75a1 100644 --- a/js/src/auto-instrumentations/index.ts +++ b/js/src/auto-instrumentations/index.ts @@ -39,6 +39,7 @@ export { openRouterConfigs } from "./configs/openrouter"; export { mistralConfigs } from "./configs/mistral"; export { googleADKConfigs } from "./configs/google-adk"; export { cohereConfigs } from "./configs/cohere"; +export { groqConfigs } from "./configs/groq"; // Re-export orchestrion configuration types // Note: ModuleMetadata and FunctionQuery are properties of InstrumentationConfig, diff --git a/js/src/exports.ts b/js/src/exports.ts index 158479847..f2d940d99 100644 --- a/js/src/exports.ts +++ b/js/src/exports.ts @@ -181,6 +181,7 @@ export { wrapOpenRouterAgent } from "./wrappers/openrouter-agent"; export { wrapOpenRouter } from "./wrappers/openrouter"; export { wrapMistral } from "./wrappers/mistral"; export { wrapCohere } from "./wrappers/cohere"; +export { wrapGroq } from "./wrappers/groq"; export { wrapVitest } from "./wrappers/vitest"; export { initNodeTestSuite } from "./wrappers/node-test"; diff --git a/js/src/instrumentation/braintrust-plugin.test.ts b/js/src/instrumentation/braintrust-plugin.test.ts index a3227c9e0..aec836b66 100644 --- a/js/src/instrumentation/braintrust-plugin.test.ts +++ b/js/src/instrumentation/braintrust-plugin.test.ts @@ -10,6 +10,7 @@ import { OpenRouterAgentPlugin } from "./plugins/openrouter-agent-plugin"; import { OpenRouterPlugin } from "./plugins/openrouter-plugin"; import { MistralPlugin } from "./plugins/mistral-plugin"; import { CoherePlugin } from "./plugins/cohere-plugin"; +import { GroqPlugin } from "./plugins/groq-plugin"; function createPluginClassMock() { return vi.fn(function MockPlugin(this: { @@ -68,6 +69,10 @@ vi.mock("./plugins/cohere-plugin", () => ({ CoherePlugin: createPluginClassMock(), })); +vi.mock("./plugins/groq-plugin", () => ({ + GroqPlugin: createPluginClassMock(), +})); + describe("BraintrustPlugin", () => { beforeEach(() => { vi.clearAllMocks(); @@ -167,6 +172,15 @@ describe("BraintrustPlugin", () => { expect(mockInstance.enable).toHaveBeenCalledTimes(1); }); + it("should create and enable Groq plugin by default", () => { + const plugin = new BraintrustPlugin(); + plugin.enable(); + + expect(GroqPlugin).toHaveBeenCalledTimes(1); + const mockInstance = vi.mocked(GroqPlugin).mock.results[0].value; + expect(mockInstance.enable).toHaveBeenCalledTimes(1); + }); + it("should create all plugins when enabled with no config", () => { const plugin = new BraintrustPlugin(); plugin.enable(); @@ -181,6 +195,7 @@ describe("BraintrustPlugin", () => { expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); expect(MistralPlugin).toHaveBeenCalledTimes(1); expect(CoherePlugin).toHaveBeenCalledTimes(1); + expect(GroqPlugin).toHaveBeenCalledTimes(1); }); it("should create all plugins when enabled with empty config", () => { @@ -197,6 +212,7 @@ describe("BraintrustPlugin", () => { expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); expect(MistralPlugin).toHaveBeenCalledTimes(1); expect(CoherePlugin).toHaveBeenCalledTimes(1); + expect(GroqPlugin).toHaveBeenCalledTimes(1); }); it("should create all plugins when enabled with empty integrations config", () => { @@ -213,6 +229,7 @@ describe("BraintrustPlugin", () => { expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); expect(MistralPlugin).toHaveBeenCalledTimes(1); expect(CoherePlugin).toHaveBeenCalledTimes(1); + expect(GroqPlugin).toHaveBeenCalledTimes(1); }); }); @@ -370,6 +387,22 @@ describe("BraintrustPlugin", () => { expect(MistralPlugin).toHaveBeenCalledTimes(1); }); + it("should not create Groq plugin when groq: false", () => { + const plugin = new BraintrustPlugin({ + integrations: { groq: false }, + }); + plugin.enable(); + + expect(GroqPlugin).not.toHaveBeenCalled(); + expect(OpenAIPlugin).toHaveBeenCalledTimes(1); + expect(AnthropicPlugin).toHaveBeenCalledTimes(1); + expect(AISDKPlugin).toHaveBeenCalledTimes(1); + expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); + expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); + expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); + expect(MistralPlugin).toHaveBeenCalledTimes(1); + }); + it("should not create OpenRouter Agent plugin when openrouterAgent: false", () => { const plugin = new BraintrustPlugin({ integrations: { openrouterAgent: false }, @@ -393,6 +426,7 @@ describe("BraintrustPlugin", () => { openrouterAgent: false, mistral: false, cohere: false, + groq: false, }, }); plugin.enable(); @@ -407,6 +441,7 @@ describe("BraintrustPlugin", () => { expect(OpenRouterAgentPlugin).not.toHaveBeenCalled(); expect(MistralPlugin).not.toHaveBeenCalled(); expect(CoherePlugin).not.toHaveBeenCalled(); + expect(GroqPlugin).not.toHaveBeenCalled(); }); it("should allow selective enabling of plugins", () => { @@ -545,6 +580,7 @@ describe("BraintrustPlugin", () => { .results[0].value; const mistralMock = vi.mocked(MistralPlugin).mock.results[0].value; const cohereMock = vi.mocked(CoherePlugin).mock.results[0].value; + const groqMock = vi.mocked(GroqPlugin).mock.results[0].value; expect(openaiMock.enable).toHaveBeenCalledTimes(1); expect(anthropicMock.enable).toHaveBeenCalledTimes(1); @@ -556,6 +592,7 @@ describe("BraintrustPlugin", () => { expect(openRouterAgentMock.enable).toHaveBeenCalledTimes(1); expect(mistralMock.enable).toHaveBeenCalledTimes(1); expect(cohereMock.enable).toHaveBeenCalledTimes(1); + expect(groqMock.enable).toHaveBeenCalledTimes(1); }); it("should disable and nullify all sub-plugins when disabled", () => { @@ -576,6 +613,7 @@ describe("BraintrustPlugin", () => { .results[0].value; const mistralMock = vi.mocked(MistralPlugin).mock.results[0].value; const cohereMock = vi.mocked(CoherePlugin).mock.results[0].value; + const groqMock = vi.mocked(GroqPlugin).mock.results[0].value; plugin.disable(); @@ -589,6 +627,7 @@ describe("BraintrustPlugin", () => { expect(openRouterAgentMock.disable).toHaveBeenCalledTimes(1); expect(mistralMock.disable).toHaveBeenCalledTimes(1); expect(cohereMock.disable).toHaveBeenCalledTimes(1); + expect(groqMock.disable).toHaveBeenCalledTimes(1); }); it("should be idempotent on multiple enable calls", () => { @@ -632,6 +671,7 @@ describe("BraintrustPlugin", () => { expect(OpenRouterAgentPlugin).not.toHaveBeenCalled(); expect(MistralPlugin).not.toHaveBeenCalled(); expect(CoherePlugin).not.toHaveBeenCalled(); + expect(GroqPlugin).not.toHaveBeenCalled(); }); it("should allow re-enabling after disable", () => { @@ -653,6 +693,7 @@ describe("BraintrustPlugin", () => { expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); expect(MistralPlugin).toHaveBeenCalledTimes(1); expect(CoherePlugin).toHaveBeenCalledTimes(1); + expect(GroqPlugin).toHaveBeenCalledTimes(1); }); it("should only disable plugins that were enabled", () => { @@ -668,6 +709,7 @@ describe("BraintrustPlugin", () => { openrouterAgent: true, mistral: false, cohere: false, + groq: true, }, }); plugin.enable(); @@ -681,6 +723,7 @@ describe("BraintrustPlugin", () => { const openRouterMock = vi.mocked(OpenRouterPlugin).mock.results[0].value; const openRouterAgentMock = vi.mocked(OpenRouterAgentPlugin).mock .results[0].value; + const groqMock = vi.mocked(GroqPlugin).mock.results[0].value; plugin.disable(); @@ -690,6 +733,7 @@ describe("BraintrustPlugin", () => { expect(huggingFaceMock.disable).toHaveBeenCalledTimes(1); expect(openRouterMock.disable).toHaveBeenCalledTimes(1); expect(openRouterAgentMock.disable).toHaveBeenCalledTimes(1); + expect(groqMock.disable).toHaveBeenCalledTimes(1); expect(MistralPlugin).not.toHaveBeenCalled(); expect(CoherePlugin).not.toHaveBeenCalled(); }); diff --git a/js/src/instrumentation/braintrust-plugin.ts b/js/src/instrumentation/braintrust-plugin.ts index 8f3919900..485d2d9bc 100644 --- a/js/src/instrumentation/braintrust-plugin.ts +++ b/js/src/instrumentation/braintrust-plugin.ts @@ -10,6 +10,7 @@ import { OpenRouterPlugin } from "./plugins/openrouter-plugin"; import { MistralPlugin } from "./plugins/mistral-plugin"; import { GoogleADKPlugin } from "./plugins/google-adk-plugin"; import { CoherePlugin } from "./plugins/cohere-plugin"; +import { GroqPlugin } from "./plugins/groq-plugin"; export interface BraintrustPluginConfig { integrations?: { @@ -26,6 +27,7 @@ export interface BraintrustPluginConfig { mistral?: boolean; googleADK?: boolean; cohere?: boolean; + groq?: boolean; }; } @@ -58,6 +60,7 @@ export class BraintrustPlugin extends BasePlugin { private mistralPlugin: MistralPlugin | null = null; private googleADKPlugin: GoogleADKPlugin | null = null; private coherePlugin: CoherePlugin | null = null; + private groqPlugin: GroqPlugin | null = null; constructor(config: BraintrustPluginConfig = {}) { super(); @@ -129,6 +132,11 @@ export class BraintrustPlugin extends BasePlugin { this.coherePlugin = new CoherePlugin(); this.coherePlugin.enable(); } + + if (integrations.groq !== false) { + this.groqPlugin = new GroqPlugin(); + this.groqPlugin.enable(); + } } protected onDisable(): void { @@ -186,6 +194,11 @@ export class BraintrustPlugin extends BasePlugin { this.coherePlugin.disable(); this.coherePlugin = null; } + + if (this.groqPlugin) { + this.groqPlugin.disable(); + this.groqPlugin = null; + } } } diff --git a/js/src/instrumentation/plugins/groq-channels.ts b/js/src/instrumentation/plugins/groq-channels.ts new file mode 100644 index 000000000..638afcbff --- /dev/null +++ b/js/src/instrumentation/plugins/groq-channels.ts @@ -0,0 +1,33 @@ +import { channel, defineChannels } from "../core/channel-definitions"; +import type { + GroqChatChoice, + GroqChatCompletion, + GroqChatCompletionChunk, + GroqChatCreateParams, + GroqChatStream, + GroqEmbeddingCreateParams, + GroqEmbeddingResponse, +} from "../../vendor-sdk-types/groq"; + +type GroqChatResult = GroqChatCompletion | GroqChatStream; + +export const groqChannels = defineChannels("groq-sdk", { + chatCompletionsCreate: channel< + [GroqChatCreateParams], + GroqChatResult, + Record, + GroqChatCompletionChunk + >({ + channelName: "chat.completions.create", + kind: "async", + }), + + embeddingsCreate: channel<[GroqEmbeddingCreateParams], GroqEmbeddingResponse>( + { + channelName: "embeddings.create", + kind: "async", + }, + ), +}); + +export type GroqChatResultChoice = GroqChatChoice; diff --git a/js/src/instrumentation/plugins/groq-plugin.test.ts b/js/src/instrumentation/plugins/groq-plugin.test.ts new file mode 100644 index 000000000..a63bab607 --- /dev/null +++ b/js/src/instrumentation/plugins/groq-plugin.test.ts @@ -0,0 +1,34 @@ +import { describe, expect, it } from "vitest"; +import { parseGroqMetrics } from "./groq-plugin"; + +describe("parseGroqMetrics", () => { + it("merges OpenAI-compatible usage metrics with Groq cache metrics", () => { + expect( + parseGroqMetrics({ + usage: { + prompt_tokens: 10, + completion_tokens: 4, + total_tokens: 14, + }, + x_groq: { + usage: { + dram_cached_tokens: 2, + sram_cached_tokens: 3, + }, + }, + }), + ).toEqual({ + completion_tokens: 4, + dram_cached_tokens: 2, + prompt_tokens: 10, + sram_cached_tokens: 3, + tokens: 14, + }); + }); + + it("returns an empty object for unknown values", () => { + expect(parseGroqMetrics(undefined)).toEqual({}); + expect(parseGroqMetrics(null)).toEqual({}); + expect(parseGroqMetrics({})).toEqual({}); + }); +}); diff --git a/js/src/instrumentation/plugins/groq-plugin.ts b/js/src/instrumentation/plugins/groq-plugin.ts new file mode 100644 index 000000000..b00ce21cc --- /dev/null +++ b/js/src/instrumentation/plugins/groq-plugin.ts @@ -0,0 +1,127 @@ +import { BasePlugin } from "../core"; +import { + traceAsyncChannel, + traceStreamingChannel, + unsubscribeAll, +} from "../core/channel-tracing"; +import { SpanTypeAttribute } from "../../../util/index"; +import { processInputAttachments } from "../../wrappers/attachment-utils"; +import { getCurrentUnixTimestamp } from "../../util"; +import { + aggregateChatCompletionChunks, + parseMetricsFromUsage, +} from "./openai-plugin"; +import { groqChannels } from "./groq-channels"; +import type { + GroqChatCompletion, + GroqChatCompletionChunk, +} from "../../vendor-sdk-types/groq"; + +export class GroqPlugin extends BasePlugin { + protected onEnable(): void { + this.unsubscribers.push( + traceStreamingChannel(groqChannels.chatCompletionsCreate, { + name: "groq.chat.completions.create", + type: SpanTypeAttribute.LLM, + extractInput: ([params]) => { + const { messages, ...metadata } = params; + return { + input: processInputAttachments(messages), + metadata: { ...metadata, provider: "groq" }, + }; + }, + extractOutput: (result) => result?.choices, + extractMetrics: (result, startTime) => { + const metrics = parseGroqMetrics(result); + if (startTime) { + metrics.time_to_first_token = getCurrentUnixTimestamp() - startTime; + } + return metrics; + }, + aggregateChunks: aggregateGroqChatCompletionChunks, + }), + ); + + this.unsubscribers.push( + traceAsyncChannel(groqChannels.embeddingsCreate, { + name: "groq.embeddings.create", + type: SpanTypeAttribute.LLM, + extractInput: ([params]) => { + const { input, ...metadata } = params; + return { + input, + metadata: { ...metadata, provider: "groq" }, + }; + }, + extractOutput: (result) => { + const embedding = result?.data?.[0]?.embedding; + return Array.isArray(embedding) + ? { embedding_length: embedding.length } + : undefined; + }, + extractMetrics: (result) => parseGroqMetrics(result), + }), + ); + } + + protected onDisable(): void { + this.unsubscribers = unsubscribeAll(this.unsubscribers); + } +} + +export function parseGroqMetrics( + result: + | Pick + | { usage?: unknown; x_groq?: unknown } + | null + | undefined, +): Record { + const metrics = parseMetricsFromUsage(result?.usage); + const xGroq = result?.x_groq; + + if (!xGroq || typeof xGroq !== "object") { + return metrics; + } + + const extraUsage = "usage" in xGroq ? xGroq.usage : undefined; + + if (!extraUsage || typeof extraUsage !== "object") { + return metrics; + } + + const dramCachedTokens = (extraUsage as Record)[ + "dram_cached_tokens" + ]; + const sramCachedTokens = (extraUsage as Record)[ + "sram_cached_tokens" + ]; + + return { + ...metrics, + ...(typeof dramCachedTokens === "number" + ? { dram_cached_tokens: dramCachedTokens } + : {}), + ...(typeof sramCachedTokens === "number" + ? { sram_cached_tokens: sramCachedTokens } + : {}), + }; +} + +export function aggregateGroqChatCompletionChunks( + chunks: GroqChatCompletionChunk[], + streamResult?: unknown, + endEvent?: unknown, +): { + metrics: Record; + output: GroqChatCompletion["choices"]; +} { + const aggregated = aggregateChatCompletionChunks( + chunks, + streamResult, + endEvent, + ); + return { + metrics: aggregated.metrics, + output: aggregated.output, + }; +} diff --git a/js/src/vendor-sdk-types/groq.ts b/js/src/vendor-sdk-types/groq.ts new file mode 100644 index 000000000..a59a72019 --- /dev/null +++ b/js/src/vendor-sdk-types/groq.ts @@ -0,0 +1,104 @@ +import type { + OpenAIAPIPromise, + OpenAIChatChoice, + OpenAIChatCompletionChunk, + OpenAIChatCreateParams, + OpenAIChatLogprobs, + OpenAIChatStream, + OpenAIEmbeddingCreateParams, + OpenAIUsage, +} from "./openai-common"; + +export interface GroqUsage extends OpenAIUsage { + queue_time?: number; +} + +export interface GroqChatCompletion { + choices: OpenAIChatChoice[]; + usage?: GroqUsage; + x_groq?: { + id?: string; + seed?: number | null; + usage?: { + dram_cached_tokens?: number; + sram_cached_tokens?: number; + [key: string]: number | undefined; + } | null; + [key: string]: unknown; + } | null; + [key: string]: unknown; +} + +export type GroqChatChoice = OpenAIChatChoice; +export type GroqChatLogprobs = OpenAIChatLogprobs; +export type GroqChatCompletionChunk = OpenAIChatCompletionChunk & { + usage?: GroqUsage; +}; +export type GroqChatCreateParams = OpenAIChatCreateParams; +export type GroqChatStream = OpenAIChatStream; + +export interface GroqEmbeddingCreateParams extends OpenAIEmbeddingCreateParams { + model?: string; +} + +export interface GroqEmbeddingResponse { + data?: Array<{ + embedding?: number[] | string; + [key: string]: unknown; + }>; + usage?: GroqUsage; + [key: string]: unknown; +} + +export interface GroqTranscriptionCreateParams { + file?: unknown; + language?: string | null; + model: string; + prompt?: string; + response_format?: string; + temperature?: number; + timestamp_granularities?: Array<"word" | "segment">; + url?: string; + [key: string]: unknown; +} + +export interface GroqTranscription { + text?: string; + [key: string]: unknown; +} + +export interface GroqChatCompletions { + create: ( + params: GroqChatCreateParams, + options?: unknown, + ) => OpenAIAPIPromise; +} + +export interface GroqChat { + completions: GroqChatCompletions; +} + +export interface GroqEmbeddings { + create: ( + params: GroqEmbeddingCreateParams, + options?: unknown, + ) => OpenAIAPIPromise; +} + +export interface GroqAudioTranscriptions { + create: ( + params: GroqTranscriptionCreateParams, + options?: unknown, + ) => OpenAIAPIPromise; +} + +export interface GroqAudio { + transcriptions: GroqAudioTranscriptions; +} + +export interface GroqClient { + audio?: GroqAudio; + chat?: GroqChat; + embeddings?: GroqEmbeddings; + [key: string]: unknown; +} diff --git a/js/src/wrappers/groq.test.ts b/js/src/wrappers/groq.test.ts new file mode 100644 index 000000000..976d727af --- /dev/null +++ b/js/src/wrappers/groq.test.ts @@ -0,0 +1,188 @@ +import { + afterEach, + beforeAll, + beforeEach, + describe, + expect, + test, + vi, +} from "vitest"; +import { configureNode } from "../node/config"; +import { _exportsForTestingOnly, initLogger } from "../logger"; +import { wrapGroq } from "./groq"; + +try { + configureNode(); +} catch { + // Best-effort initialization for test environments. +} + +describe("groq wrapper", () => { + let backgroundLogger: ReturnType< + typeof _exportsForTestingOnly.useTestBackgroundLogger + >; + + beforeAll(async () => { + await _exportsForTestingOnly.simulateLoginForTests(); + }); + + beforeEach(() => { + backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); + initLogger({ + projectId: "test-project-id", + projectName: "groq.test.ts", + }); + }); + + afterEach(() => { + _exportsForTestingOnly.clearTestBackgroundLogger(); + vi.restoreAllMocks(); + }); + + test("returns original object for unsupported clients", () => { + const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); + const invalid = { foo: "bar" }; + + expect(wrapGroq(invalid)).toBe(invalid); + expect(warnSpy).toHaveBeenCalledWith( + "Unsupported Groq library. Not wrapping.", + ); + }); + + test("wraps chat completions and embeddings", async () => { + async function* stream() { + yield { + choices: [{ delta: { role: "assistant" }, finish_reason: null }], + }; + yield { + choices: [ + { + delta: { content: "STREAM" }, + finish_reason: "stop", + }, + ], + usage: { + completion_tokens: 1, + prompt_tokens: 4, + total_tokens: 5, + }, + }; + } + + const wrapped = wrapGroq({ + chat: { + completions: { + create: vi.fn(async (request: Record) => { + if (request.stream) { + return stream(); + } + + return { + choices: [ + { + index: 0, + message: { + content: "OK", + role: "assistant", + }, + }, + ], + usage: { + completion_tokens: 2, + prompt_tokens: 5, + total_tokens: 7, + }, + x_groq: { + usage: { + dram_cached_tokens: 1, + sram_cached_tokens: 2, + }, + }, + }; + }), + }, + }, + embeddings: { + create: vi.fn(async () => ({ + data: [{ embedding: [0.1, 0.2, 0.3] }], + usage: { + prompt_tokens: 3, + total_tokens: 3, + }, + })), + }, + withOptions(options: unknown) { + return options ? this : null; + }, + }); + + expect(wrapped.withOptions({})).toBe(wrapped); + + await wrapped.chat.completions.create({ + max_completion_tokens: 12, + messages: [{ content: "Reply with exactly OK.", role: "user" }], + model: "llama-3.3-70b-versatile", + temperature: 0, + }); + + const streamed = await wrapped.chat.completions.create({ + messages: [{ content: "Reply with exactly STREAM.", role: "user" }], + model: "llama-3.3-70b-versatile", + stream: true, + }); + for await (const _chunk of streamed) { + // Consume the stream so chunk aggregation runs. + } + + await wrapped.embeddings.create({ + input: "braintrust tracing", + model: "nomic-embed-text-v1_5", + }); + + const spans = await backgroundLogger.drain(); + expect(spans).toHaveLength(3); + + const chatSpan = spans.find( + (span: any) => + span.span_attributes?.name === "groq.chat.completions.create" && + span.output?.[0]?.message?.content === "OK", + ) as Record | undefined; + const streamSpan = spans.find( + (span: any) => + span.span_attributes?.name === "groq.chat.completions.create" && + span.output?.[0]?.message?.content === "STREAM", + ) as Record | undefined; + const embeddingSpan = spans.find( + (span: any) => span.span_attributes?.name === "groq.embeddings.create", + ) as Record | undefined; + + expect(chatSpan?.metadata).toMatchObject({ + model: "llama-3.3-70b-versatile", + provider: "groq", + temperature: 0, + }); + expect(chatSpan?.metrics).toMatchObject({ + completion_tokens: 2, + dram_cached_tokens: 1, + prompt_tokens: 5, + sram_cached_tokens: 2, + time_to_first_token: expect.any(Number), + tokens: 7, + }); + + expect(streamSpan?.metrics).toMatchObject({ + completion_tokens: 1, + prompt_tokens: 4, + time_to_first_token: expect.any(Number), + tokens: 5, + }); + + expect(embeddingSpan?.metadata).toMatchObject({ + model: "nomic-embed-text-v1_5", + provider: "groq", + }); + expect(embeddingSpan?.output).toEqual({ + embedding_length: 3, + }); + }); +}); diff --git a/js/src/wrappers/groq.ts b/js/src/wrappers/groq.ts new file mode 100644 index 000000000..68312fa81 --- /dev/null +++ b/js/src/wrappers/groq.ts @@ -0,0 +1,159 @@ +import { groqChannels } from "../instrumentation/plugins/groq-channels"; +import type { + GroqChat, + GroqChatCompletion, + GroqChatCreateParams, + GroqChatStream, + GroqClient, + GroqEmbeddingCreateParams, + GroqEmbeddingResponse, + GroqEmbeddings, +} from "../vendor-sdk-types/groq"; + +/** + * Wrap a Groq client (created with `new Groq(...)`) with Braintrust tracing. + */ +export function wrapGroq(groq: T): T { + if (isSupportedGroqClient(groq)) { + return groqProxy(groq) as T; + } + + // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. + console.warn("Unsupported Groq library. Not wrapping."); + return groq; +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null; +} + +function hasFunction(value: unknown, methodName: string): boolean { + return ( + isRecord(value) && + methodName in value && + typeof value[methodName] === "function" + ); +} + +function hasChat(value: unknown): value is GroqChat { + return ( + isRecord(value) && + isRecord(value.completions) && + hasFunction(value.completions, "create") + ); +} + +function hasEmbeddings(value: unknown): value is GroqEmbeddings { + return hasFunction(value, "create"); +} + +function isSupportedGroqClient(value: unknown): value is GroqClient { + return ( + isRecord(value) && + ((value.chat !== undefined && hasChat(value.chat)) || + (value.embeddings !== undefined && hasEmbeddings(value.embeddings))) + ); +} + +function groqProxy(groq: GroqClient): GroqClient { + const privateMethodWorkaroundCache = new WeakMap< + (...args: unknown[]) => unknown, + (...args: unknown[]) => unknown + >(); + + const completionProxy = groq.chat?.completions + ? new Proxy(groq.chat.completions, { + get(target, prop, receiver) { + if (prop === "create") { + return wrapChatCompletionsCreate(target.create.bind(target)); + } + + return Reflect.get(target, prop, receiver); + }, + }) + : undefined; + + const chatProxy = groq.chat + ? new Proxy(groq.chat, { + get(target, prop, receiver) { + if (prop === "completions") { + return completionProxy ?? target.completions; + } + + return Reflect.get(target, prop, receiver); + }, + }) + : undefined; + + const embeddingsProxy = groq.embeddings + ? new Proxy(groq.embeddings, { + get(target, prop, receiver) { + if (prop === "create") { + return wrapEmbeddingsCreate(target.create.bind(target)); + } + + return Reflect.get(target, prop, receiver); + }, + }) + : undefined; + + const topLevelProxy: GroqClient = new Proxy(groq, { + get(target, prop, receiver) { + switch (prop) { + case "chat": + return chatProxy ?? target.chat; + case "embeddings": + return embeddingsProxy ?? target.embeddings; + } + + const value = Reflect.get(target, prop, target); + if (typeof value !== "function") { + return value; + } + + const cachedValue = privateMethodWorkaroundCache.get(value); + if (cachedValue) { + return cachedValue; + } + + const thisBoundValue = function ( + this: unknown, + ...args: unknown[] + ): unknown { + const thisArg = this === topLevelProxy ? target : this; + const output = Reflect.apply(value, thisArg, args); + return output === target ? topLevelProxy : output; + }; + + privateMethodWorkaroundCache.set(value, thisBoundValue); + return thisBoundValue; + }, + }); + + return topLevelProxy; +} + +function wrapChatCompletionsCreate( + create: ( + request: GroqChatCreateParams, + options?: unknown, + ) => Promise, +): GroqChat["completions"]["create"] { + return (request, options) => + groqChannels.chatCompletionsCreate.tracePromise( + () => create(request, options), + { arguments: [request] }, + ) as ReturnType; +} + +function wrapEmbeddingsCreate( + create: ( + request: GroqEmbeddingCreateParams, + options?: unknown, + ) => Promise, +): GroqEmbeddings["create"] { + return (request, options) => + groqChannels.embeddingsCreate.tracePromise(() => create(request, options), { + arguments: [request], + }) as ReturnType; +} diff --git a/turbo.json b/turbo.json index 22c028ca6..bad5b66ae 100644 --- a/turbo.json +++ b/turbo.json @@ -6,6 +6,7 @@ "ANTHROPIC_API_KEY", "GEMINI_API_KEY", "COHERE_API_KEY", + "GROQ_API_KEY", "OPENROUTER_API_KEY", "MISTRAL_API_KEY", "HUGGINGFACE_API_KEY" @@ -26,6 +27,7 @@ "BRAINTRUST_API_KEY", "GEMINI_API_KEY", "COHERE_API_KEY", + "GROQ_API_KEY", "OPENAI_API_KEY", "OPENAI_BASE_URL", "OPENROUTER_API_KEY", @@ -44,6 +46,7 @@ "BRAINTRUST_E2E_RUN_CONTEXT_DIR", "GEMINI_API_KEY", "COHERE_API_KEY", + "GROQ_API_KEY", "OPENAI_API_KEY", "OPENAI_BASE_URL", "OPENROUTER_API_KEY", @@ -67,6 +70,7 @@ "BRAINTRUST_E2E_RUN_CONTEXT_DIR", "GEMINI_API_KEY", "COHERE_API_KEY", + "GROQ_API_KEY", "OPENAI_API_KEY", "OPENAI_BASE_URL", "OPENROUTER_API_KEY", @@ -85,6 +89,7 @@ "BRAINTRUST_E2E_RUN_CONTEXT_DIR", "GEMINI_API_KEY", "COHERE_API_KEY", + "GROQ_API_KEY", "OPENAI_API_KEY", "OPENAI_BASE_URL", "OPENROUTER_API_KEY", @@ -117,6 +122,7 @@ "BRAINTRUST_API_KEY", "GEMINI_API_KEY", "COHERE_API_KEY", + "GROQ_API_KEY", "OPENAI_API_KEY", "OPENROUTER_API_KEY", "MISTRAL_API_KEY", @@ -131,6 +137,7 @@ "BRAINTRUST_API_KEY", "GEMINI_API_KEY", "COHERE_API_KEY", + "GROQ_API_KEY", "OPENAI_API_KEY", "OPENROUTER_API_KEY", "MISTRAL_API_KEY", @@ -145,6 +152,7 @@ "BRAINTRUST_API_KEY", "GEMINI_API_KEY", "COHERE_API_KEY", + "GROQ_API_KEY", "OPENAI_API_KEY", "OPENROUTER_API_KEY", "MISTRAL_API_KEY", From c8bfbffceee9a982528792020f29e2dc074ff8e4 Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Tue, 21 Apr 2026 17:11:53 +0200 Subject: [PATCH 2/4] changeset --- .changeset/fresh-crabs-dream.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .changeset/fresh-crabs-dream.md diff --git a/.changeset/fresh-crabs-dream.md b/.changeset/fresh-crabs-dream.md new file mode 100644 index 000000000..4f35ee20e --- /dev/null +++ b/.changeset/fresh-crabs-dream.md @@ -0,0 +1,5 @@ +--- +"braintrust": patch +--- + +feat: Add instrumentation for groq-sdk From b2a25f4c0574890f62276e43a6f3dfaae5f3e75e Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Tue, 21 Apr 2026 17:14:44 +0200 Subject: [PATCH 3/4] dead code --- js/src/instrumentation/plugins/groq-plugin.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/js/src/instrumentation/plugins/groq-plugin.ts b/js/src/instrumentation/plugins/groq-plugin.ts index b00ce21cc..777bc40a0 100644 --- a/js/src/instrumentation/plugins/groq-plugin.ts +++ b/js/src/instrumentation/plugins/groq-plugin.ts @@ -107,7 +107,7 @@ export function parseGroqMetrics( }; } -export function aggregateGroqChatCompletionChunks( +function aggregateGroqChatCompletionChunks( chunks: GroqChatCompletionChunk[], streamResult?: unknown, endEvent?: unknown, From 5bbe618c0f53f44fd9178043d9aea283bfe32598 Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Tue, 21 Apr 2026 17:38:54 +0200 Subject: [PATCH 4/4] summary --- e2e/config/pr-comment-scenarios.json | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/e2e/config/pr-comment-scenarios.json b/e2e/config/pr-comment-scenarios.json index 7870cbcfe..4820d7d24 100644 --- a/e2e/config/pr-comment-scenarios.json +++ b/e2e/config/pr-comment-scenarios.json @@ -39,6 +39,15 @@ { "variantKey": "google-genai-v1460", "label": "v1.46.0" } ] }, + { + "scenarioDirName": "groq-instrumentation", + "label": "Groq Instrumentation", + "metadataScenario": "groq-instrumentation", + "variants": [ + { "variantKey": "groq-v1-wrapped", "label": "Wrapped" }, + { "variantKey": "groq-v1-auto", "label": "Auto-hook" } + ] + }, { "scenarioDirName": "huggingface-instrumentation", "label": "HuggingFace Instrumentation",