From f23bd4d0fedb225d0de28a434b8d991444691647 Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Thu, 20 Aug 2026 14:13:39 +0000 Subject: [PATCH 1/2] wip feat: Add instrumentation API for groq batch API --- .../__snapshots__/groq-v1-auto.span-tree.json | 227 +- .../__snapshots__/groq-v1-auto.span-tree.txt | 304 +- .../groq-v1-latest-auto.span-tree.json | 227 +- .../groq-v1-latest-auto.span-tree.txt | 304 +- .../groq-v1-latest-wrapped.span-tree.json | 227 +- .../groq-v1-latest-wrapped.span-tree.txt | 304 +- .../groq-v1-wrapped.span-tree.json | 227 +- .../groq-v1-wrapped.span-tree.txt | 304 +- .../groq-instrumentation/assertions.ts | 123 + .../groq-instrumentation/scenario.impl.mjs | 224 +- js/src/exports.ts | 21 + js/src/groq-batch-types.ts | 82 + js/src/groq-batch.ts | 116 + .../groq-batch-instrumentation.test.ts | 1411 +++++++++ .../plugins/groq-batch-instrumentation.ts | 2786 +++++++++++++++++ .../instrumentation/plugins/groq-channels.ts | 38 + .../plugins/groq-plugin.test.ts | 202 +- js/src/instrumentation/plugins/groq-plugin.ts | 84 +- .../instrumentation/plugins/groq-span-data.ts | 175 ++ js/src/logger.ts | 33 +- js/src/wrappers/groq.test.ts | 3 +- 21 files changed, 6892 insertions(+), 530 deletions(-) create mode 100644 js/src/groq-batch-types.ts create mode 100644 js/src/groq-batch.ts create mode 100644 js/src/instrumentation/plugins/groq-batch-instrumentation.test.ts create mode 100644 js/src/instrumentation/plugins/groq-batch-instrumentation.ts create mode 100644 js/src/instrumentation/plugins/groq-span-data.ts diff --git a/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-auto.span-tree.json b/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-auto.span-tree.json index 84ba50ea9..8b29ec775 100644 --- a/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-auto.span-tree.json +++ b/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-auto.span-tree.json @@ -35,14 +35,10 @@ "temperature": 0 }, "metrics": { - "completion_time": 0.0072439, "completion_tokens": 2, - "prompt_time": 0.001728151, "prompt_tokens": 40, - "queue_time": 0.054452998, "time_to_first_token": 0, - "tokens": 42, - "total_time": 0.008972051 + "tokens": 42 } } ], @@ -78,18 +74,13 @@ "metadata": { "model": "llama-3.3-70b-versatile", "provider": "groq", - "stream": true, "temperature": 0 }, "metrics": { - "completion_time": 0.006718869, "completion_tokens": 2, - "prompt_time": 0.003972294, "prompt_tokens": 40, - "queue_time": 0.095351581, "time_to_first_token": 0, - "tokens": 42, - "total_time": 0.010691163 + "tokens": 42 } } ], @@ -128,19 +119,14 @@ "model": "openai/gpt-oss-20b", "provider": "groq", "reasoning_format": "parsed", - "stream": true, "temperature": 0.6 }, "metrics": { "completion_reasoning_tokens": 158, - "completion_time": 0.270633239, "completion_tokens": 168, - "prompt_time": 0.005148062, "prompt_tokens": 108, - "queue_time": 0.103912089, "time_to_first_token": 0, - "tokens": 276, - "total_time": 0.275781301 + "tokens": 276 } } ], @@ -215,14 +201,10 @@ ] }, "metrics": { - "completion_time": 0.039504246, "completion_tokens": 11, - "prompt_time": 0.020812356, "prompt_tokens": 246, - "queue_time": 0.091778475, "time_to_first_token": 0, - "tokens": 257, - "total_time": 0.060316602 + "tokens": 257 } } ], @@ -230,6 +212,207 @@ "operation": "tool", "testRunId": "" } + }, + { + "name": "groq-batch-operation", + "children": [ + { + "name": "groq.batch", + "type": "task", + "children": [ + { + "name": "groq.chat.completions.create", + "type": "llm", + "children": [], + "input": [ + { + "content": "Reply with ALPHA.", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "ALPHA", + "role": "assistant" + } + } + ], + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "metrics": { + "completion_tokens": 1, + "prompt_cached_tokens": 3, + "prompt_tokens": 5, + "tokens": 6 + } + }, + { + "name": "groq.chat.completions.create", + "type": "llm", + "children": [], + "input": [ + { + "content": "Reply with BRAVO.", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "BRAVO", + "role": "assistant" + } + } + ], + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "metrics": { + "completion_tokens": 1, + "prompt_cached_tokens": 3, + "prompt_tokens": 5, + "tokens": 6 + } + }, + { + "name": "groq.chat.completions.create", + "type": "llm", + "children": [], + "input": [ + { + "content": "This request should fail.", + "role": "user" + } + ], + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "error": "Groq batch fixture request failed" + } + ], + "metadata": { + "provider": "groq" + } + } + ], + "metadata": { + "operation": "batch", + "testRunId": "" + } + }, + { + "name": "groq-batch-collect-only-operation", + "children": [ + { + "name": "groq.batch", + "type": "task", + "children": [ + { + "name": "groq.chat.completions.create", + "type": "llm", + "children": [], + "input": [ + { + "content": "Collect this response.", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "COLLECTED", + "role": "assistant" + } + } + ], + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "metrics": { + "completion_tokens": 1, + "prompt_tokens": 4, + "tokens": 5 + } + } + ], + "metadata": { + "provider": "groq" + } + } + ], + "metadata": { + "operation": "batch-collect-only", + "testRunId": "" + } + }, + { + "name": "groq-batch-submission-failure-operation", + "children": [ + { + "name": "groq.batch", + "type": "task", + "children": [ + { + "name": "groq.chat.completions.create", + "type": "llm", + "children": [], + "input": [ + { + "content": "First pending request.", + "role": "user" + } + ], + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "error": "Groq batch submission failed" + }, + { + "name": "groq.chat.completions.create", + "type": "llm", + "children": [], + "input": [ + { + "content": "Second pending request.", + "role": "user" + } + ], + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "error": "Groq batch submission failed" + } + ], + "metadata": { + "provider": "groq" + }, + "error": "Groq batch submission failed" + } + ], + "metadata": { + "operation": "batch-submission-failure", + "testRunId": "" + } } ], "metadata": { diff --git a/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-auto.span-tree.txt b/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-auto.span-tree.txt index 8bab80353..dfc41ece9 100644 --- a/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-auto.span-tree.txt +++ b/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-auto.span-tree.txt @@ -34,14 +34,10 @@ span_tree: │ "temperature": 0 │ } │ metrics: { - │ "completion_time": 0.0072439, │ "completion_tokens": 2, - │ "prompt_time": 0.001728151, │ "prompt_tokens": 40, - │ "queue_time": 0.054452998, │ "time_to_first_token": 0, - │ "tokens": 42, - │ "total_time": 0.008972051 + │ "tokens": 42 │ } ├── groq-stream-operation │ metadata: { @@ -69,18 +65,13 @@ span_tree: │ metadata: { │ "model": "llama-3.3-70b-versatile", │ "provider": "groq", - │ "stream": true, │ "temperature": 0 │ } │ metrics: { - │ "completion_time": 0.006718869, │ "completion_tokens": 2, - │ "prompt_time": 0.003972294, │ "prompt_tokens": 40, - │ "queue_time": 0.095351581, │ "time_to_first_token": 0, - │ "tokens": 42, - │ "total_time": 0.010691163 + │ "tokens": 42 │ } ├── groq-reasoning-stream-operation │ metadata: { @@ -111,91 +102,232 @@ span_tree: │ "model": "openai/gpt-oss-20b", │ "provider": "groq", │ "reasoning_format": "parsed", - │ "stream": true, │ "temperature": 0.6 │ } │ metrics: { │ "completion_reasoning_tokens": 158, - │ "completion_time": 0.270633239, │ "completion_tokens": 168, - │ "prompt_time": 0.005148062, │ "prompt_tokens": 108, - │ "queue_time": 0.103912089, │ "time_to_first_token": 0, - │ "tokens": 276, - │ "total_time": 0.275781301 + │ "tokens": 276 │ } - └── groq-tool-operation + ├── groq-tool-operation + │ metadata: { + │ "operation": "tool", + │ "testRunId": "" + │ } + │ └── groq.chat.completions.create [llm] + │ input: [ + │ { + │ "content": "Check the weather in Vienna and use the weather tool.", + │ "role": "user" + │ } + │ ] + │ output: [ + │ { + │ "finish_reason": "tool_calls", + │ "index": 0, + │ "logprobs": null, + │ "message": { + │ "role": "assistant", + │ "tool_calls": [ + │ { + │ "function": { + │ "arguments": "{\"location\":\"Vienna\"}", + │ "name": "get_weather" + │ }, + │ "id": "", + │ "type": "function" + │ } + │ ] + │ } + │ } + │ ] + │ metadata: { + │ "model": "llama-3.3-70b-versatile", + │ "provider": "groq", + │ "temperature": 0, + │ "tool_choice": { + │ "function": { + │ "name": "get_weather" + │ }, + │ "type": "function" + │ }, + │ "tools": [ + │ { + │ "function": { + │ "description": "Get the weather for a city.", + │ "name": "get_weather", + │ "parameters": { + │ "properties": { + │ "location": { + │ "description": "City name.", + │ "type": "string" + │ } + │ }, + │ "required": [ + │ "location" + │ ], + │ "type": "object" + │ } + │ }, + │ "type": "function" + │ } + │ ] + │ } + │ metrics: { + │ "completion_tokens": 11, + │ "prompt_tokens": 246, + │ "time_to_first_token": 0, + │ "tokens": 257 + │ } + ├── groq-batch-operation + │ metadata: { + │ "operation": "batch", + │ "testRunId": "" + │ } + │ └── groq.batch [task] + │ metadata: { + │ "provider": "groq" + │ } + │ ├── groq.chat.completions.create [llm] + │ │ input: [ + │ │ { + │ │ "content": "Reply with ALPHA.", + │ │ "role": "user" + │ │ } + │ │ ] + │ │ output: [ + │ │ { + │ │ "finish_reason": "stop", + │ │ "index": 0, + │ │ "message": { + │ │ "content": "ALPHA", + │ │ "role": "assistant" + │ │ } + │ │ } + │ │ ] + │ │ metadata: { + │ │ "model": "llama-3.3-70b-versatile", + │ │ "provider": "groq", + │ │ "temperature": 0 + │ │ } + │ │ metrics: { + │ │ "completion_tokens": 1, + │ │ "prompt_cached_tokens": 3, + │ │ "prompt_tokens": 5, + │ │ "tokens": 6 + │ │ } + │ ├── groq.chat.completions.create [llm] + │ │ input: [ + │ │ { + │ │ "content": "Reply with BRAVO.", + │ │ "role": "user" + │ │ } + │ │ ] + │ │ output: [ + │ │ { + │ │ "finish_reason": "stop", + │ │ "index": 0, + │ │ "message": { + │ │ "content": "BRAVO", + │ │ "role": "assistant" + │ │ } + │ │ } + │ │ ] + │ │ metadata: { + │ │ "model": "llama-3.3-70b-versatile", + │ │ "provider": "groq", + │ │ "temperature": 0 + │ │ } + │ │ metrics: { + │ │ "completion_tokens": 1, + │ │ "prompt_cached_tokens": 3, + │ │ "prompt_tokens": 5, + │ │ "tokens": 6 + │ │ } + │ └── groq.chat.completions.create [llm] + │ input: [ + │ { + │ "content": "This request should fail.", + │ "role": "user" + │ } + │ ] + │ metadata: { + │ "model": "llama-3.3-70b-versatile", + │ "provider": "groq", + │ "temperature": 0 + │ } + │ error: "Groq batch fixture request failed" + ├── groq-batch-collect-only-operation + │ metadata: { + │ "operation": "batch-collect-only", + │ "testRunId": "" + │ } + │ └── groq.batch [task] + │ metadata: { + │ "provider": "groq" + │ } + │ └── groq.chat.completions.create [llm] + │ input: [ + │ { + │ "content": "Collect this response.", + │ "role": "user" + │ } + │ ] + │ output: [ + │ { + │ "finish_reason": "stop", + │ "index": 0, + │ "message": { + │ "content": "COLLECTED", + │ "role": "assistant" + │ } + │ } + │ ] + │ metadata: { + │ "model": "llama-3.3-70b-versatile", + │ "provider": "groq", + │ "temperature": 0 + │ } + │ metrics: { + │ "completion_tokens": 1, + │ "prompt_tokens": 4, + │ "tokens": 5 + │ } + └── groq-batch-submission-failure-operation metadata: { - "operation": "tool", + "operation": "batch-submission-failure", "testRunId": "" } - └── groq.chat.completions.create [llm] - input: [ - { - "content": "Check the weather in Vienna and use the weather tool.", - "role": "user" - } - ] - output: [ - { - "finish_reason": "tool_calls", - "index": 0, - "logprobs": null, - "message": { - "role": "assistant", - "tool_calls": [ - { - "function": { - "arguments": "{\"location\":\"Vienna\"}", - "name": "get_weather" - }, - "id": "", - "type": "function" - } - ] - } - } - ] + └── groq.batch [task] metadata: { - "model": "llama-3.3-70b-versatile", - "provider": "groq", - "temperature": 0, - "tool_choice": { - "function": { - "name": "get_weather" - }, - "type": "function" - }, - "tools": [ - { - "function": { - "description": "Get the weather for a city.", - "name": "get_weather", - "parameters": { - "properties": { - "location": { - "description": "City name.", - "type": "string" - } - }, - "required": [ - "location" - ], - "type": "object" - } - }, - "type": "function" - } - ] - } - metrics: { - "completion_time": 0.039504246, - "completion_tokens": 11, - "prompt_time": 0.020812356, - "prompt_tokens": 246, - "queue_time": 0.091778475, - "time_to_first_token": 0, - "tokens": 257, - "total_time": 0.060316602 + "provider": "groq" } + error: "Groq batch submission failed" + ├── groq.chat.completions.create [llm] + │ input: [ + │ { + │ "content": "First pending request.", + │ "role": "user" + │ } + │ ] + │ metadata: { + │ "model": "llama-3.3-70b-versatile", + │ "provider": "groq", + │ "temperature": 0 + │ } + │ error: "Groq batch submission failed" + └── groq.chat.completions.create [llm] + input: [ + { + "content": "Second pending request.", + "role": "user" + } + ] + metadata: { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + } + error: "Groq batch submission failed" diff --git a/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-latest-auto.span-tree.json b/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-latest-auto.span-tree.json index 01ec3abcd..3affda15a 100644 --- a/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-latest-auto.span-tree.json +++ b/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-latest-auto.span-tree.json @@ -35,14 +35,10 @@ "temperature": 0 }, "metrics": { - "completion_time": 0.007446148, "completion_tokens": 2, - "prompt_time": 0.001747341, "prompt_tokens": 40, - "queue_time": 0.036266065, "time_to_first_token": 0, - "tokens": 42, - "total_time": 0.009193489 + "tokens": 42 } } ], @@ -78,18 +74,13 @@ "metadata": { "model": "llama-3.3-70b-versatile", "provider": "groq", - "stream": true, "temperature": 0 }, "metrics": { - "completion_time": 0.008458721, "completion_tokens": 2, - "prompt_time": 0.001661312, "prompt_tokens": 40, - "queue_time": 0.036348217, "time_to_first_token": 0, - "tokens": 42, - "total_time": 0.010120033 + "tokens": 42 } } ], @@ -128,19 +119,14 @@ "model": "openai/gpt-oss-20b", "provider": "groq", "reasoning_format": "parsed", - "stream": true, "temperature": 0.6 }, "metrics": { "completion_reasoning_tokens": 173, - "completion_time": 0.185192367, "completion_tokens": 183, - "prompt_time": 0.005113732, "prompt_tokens": 108, - "queue_time": 0.07155807, "time_to_first_token": 0, - "tokens": 291, - "total_time": 0.190306099 + "tokens": 291 } } ], @@ -215,14 +201,10 @@ ] }, "metrics": { - "completion_time": 0.035807806, "completion_tokens": 11, - "prompt_time": 0.025119244, "prompt_tokens": 246, - "queue_time": 0.145082373, "time_to_first_token": 0, - "tokens": 257, - "total_time": 0.06092705 + "tokens": 257 } } ], @@ -230,6 +212,207 @@ "operation": "tool", "testRunId": "" } + }, + { + "name": "groq-batch-operation", + "children": [ + { + "name": "groq.batch", + "type": "task", + "children": [ + { + "name": "groq.chat.completions.create", + "type": "llm", + "children": [], + "input": [ + { + "content": "Reply with ALPHA.", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "ALPHA", + "role": "assistant" + } + } + ], + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "metrics": { + "completion_tokens": 1, + "prompt_cached_tokens": 3, + "prompt_tokens": 5, + "tokens": 6 + } + }, + { + "name": "groq.chat.completions.create", + "type": "llm", + "children": [], + "input": [ + { + "content": "Reply with BRAVO.", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "BRAVO", + "role": "assistant" + } + } + ], + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "metrics": { + "completion_tokens": 1, + "prompt_cached_tokens": 3, + "prompt_tokens": 5, + "tokens": 6 + } + }, + { + "name": "groq.chat.completions.create", + "type": "llm", + "children": [], + "input": [ + { + "content": "This request should fail.", + "role": "user" + } + ], + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "error": "Groq batch fixture request failed" + } + ], + "metadata": { + "provider": "groq" + } + } + ], + "metadata": { + "operation": "batch", + "testRunId": "" + } + }, + { + "name": "groq-batch-collect-only-operation", + "children": [ + { + "name": "groq.batch", + "type": "task", + "children": [ + { + "name": "groq.chat.completions.create", + "type": "llm", + "children": [], + "input": [ + { + "content": "Collect this response.", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "COLLECTED", + "role": "assistant" + } + } + ], + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "metrics": { + "completion_tokens": 1, + "prompt_tokens": 4, + "tokens": 5 + } + } + ], + "metadata": { + "provider": "groq" + } + } + ], + "metadata": { + "operation": "batch-collect-only", + "testRunId": "" + } + }, + { + "name": "groq-batch-submission-failure-operation", + "children": [ + { + "name": "groq.batch", + "type": "task", + "children": [ + { + "name": "groq.chat.completions.create", + "type": "llm", + "children": [], + "input": [ + { + "content": "First pending request.", + "role": "user" + } + ], + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "error": "Groq batch submission failed" + }, + { + "name": "groq.chat.completions.create", + "type": "llm", + "children": [], + "input": [ + { + "content": "Second pending request.", + "role": "user" + } + ], + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "error": "Groq batch submission failed" + } + ], + "metadata": { + "provider": "groq" + }, + "error": "Groq batch submission failed" + } + ], + "metadata": { + "operation": "batch-submission-failure", + "testRunId": "" + } } ], "metadata": { diff --git a/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-latest-auto.span-tree.txt b/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-latest-auto.span-tree.txt index 215dd4248..3a3ff791e 100644 --- a/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-latest-auto.span-tree.txt +++ b/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-latest-auto.span-tree.txt @@ -34,14 +34,10 @@ span_tree: │ "temperature": 0 │ } │ metrics: { - │ "completion_time": 0.007446148, │ "completion_tokens": 2, - │ "prompt_time": 0.001747341, │ "prompt_tokens": 40, - │ "queue_time": 0.036266065, │ "time_to_first_token": 0, - │ "tokens": 42, - │ "total_time": 0.009193489 + │ "tokens": 42 │ } ├── groq-stream-operation │ metadata: { @@ -69,18 +65,13 @@ span_tree: │ metadata: { │ "model": "llama-3.3-70b-versatile", │ "provider": "groq", - │ "stream": true, │ "temperature": 0 │ } │ metrics: { - │ "completion_time": 0.008458721, │ "completion_tokens": 2, - │ "prompt_time": 0.001661312, │ "prompt_tokens": 40, - │ "queue_time": 0.036348217, │ "time_to_first_token": 0, - │ "tokens": 42, - │ "total_time": 0.010120033 + │ "tokens": 42 │ } ├── groq-reasoning-stream-operation │ metadata: { @@ -111,91 +102,232 @@ span_tree: │ "model": "openai/gpt-oss-20b", │ "provider": "groq", │ "reasoning_format": "parsed", - │ "stream": true, │ "temperature": 0.6 │ } │ metrics: { │ "completion_reasoning_tokens": 173, - │ "completion_time": 0.185192367, │ "completion_tokens": 183, - │ "prompt_time": 0.005113732, │ "prompt_tokens": 108, - │ "queue_time": 0.07155807, │ "time_to_first_token": 0, - │ "tokens": 291, - │ "total_time": 0.190306099 + │ "tokens": 291 │ } - └── groq-tool-operation + ├── groq-tool-operation + │ metadata: { + │ "operation": "tool", + │ "testRunId": "" + │ } + │ └── groq.chat.completions.create [llm] + │ input: [ + │ { + │ "content": "Check the weather in Vienna and use the weather tool.", + │ "role": "user" + │ } + │ ] + │ output: [ + │ { + │ "finish_reason": "tool_calls", + │ "index": 0, + │ "logprobs": null, + │ "message": { + │ "role": "assistant", + │ "tool_calls": [ + │ { + │ "function": { + │ "arguments": "{\"location\":\"Vienna\"}", + │ "name": "get_weather" + │ }, + │ "id": "", + │ "type": "function" + │ } + │ ] + │ } + │ } + │ ] + │ metadata: { + │ "model": "llama-3.3-70b-versatile", + │ "provider": "groq", + │ "temperature": 0, + │ "tool_choice": { + │ "function": { + │ "name": "get_weather" + │ }, + │ "type": "function" + │ }, + │ "tools": [ + │ { + │ "function": { + │ "description": "Get the weather for a city.", + │ "name": "get_weather", + │ "parameters": { + │ "properties": { + │ "location": { + │ "description": "City name.", + │ "type": "string" + │ } + │ }, + │ "required": [ + │ "location" + │ ], + │ "type": "object" + │ } + │ }, + │ "type": "function" + │ } + │ ] + │ } + │ metrics: { + │ "completion_tokens": 11, + │ "prompt_tokens": 246, + │ "time_to_first_token": 0, + │ "tokens": 257 + │ } + ├── groq-batch-operation + │ metadata: { + │ "operation": "batch", + │ "testRunId": "" + │ } + │ └── groq.batch [task] + │ metadata: { + │ "provider": "groq" + │ } + │ ├── groq.chat.completions.create [llm] + │ │ input: [ + │ │ { + │ │ "content": "Reply with ALPHA.", + │ │ "role": "user" + │ │ } + │ │ ] + │ │ output: [ + │ │ { + │ │ "finish_reason": "stop", + │ │ "index": 0, + │ │ "message": { + │ │ "content": "ALPHA", + │ │ "role": "assistant" + │ │ } + │ │ } + │ │ ] + │ │ metadata: { + │ │ "model": "llama-3.3-70b-versatile", + │ │ "provider": "groq", + │ │ "temperature": 0 + │ │ } + │ │ metrics: { + │ │ "completion_tokens": 1, + │ │ "prompt_cached_tokens": 3, + │ │ "prompt_tokens": 5, + │ │ "tokens": 6 + │ │ } + │ ├── groq.chat.completions.create [llm] + │ │ input: [ + │ │ { + │ │ "content": "Reply with BRAVO.", + │ │ "role": "user" + │ │ } + │ │ ] + │ │ output: [ + │ │ { + │ │ "finish_reason": "stop", + │ │ "index": 0, + │ │ "message": { + │ │ "content": "BRAVO", + │ │ "role": "assistant" + │ │ } + │ │ } + │ │ ] + │ │ metadata: { + │ │ "model": "llama-3.3-70b-versatile", + │ │ "provider": "groq", + │ │ "temperature": 0 + │ │ } + │ │ metrics: { + │ │ "completion_tokens": 1, + │ │ "prompt_cached_tokens": 3, + │ │ "prompt_tokens": 5, + │ │ "tokens": 6 + │ │ } + │ └── groq.chat.completions.create [llm] + │ input: [ + │ { + │ "content": "This request should fail.", + │ "role": "user" + │ } + │ ] + │ metadata: { + │ "model": "llama-3.3-70b-versatile", + │ "provider": "groq", + │ "temperature": 0 + │ } + │ error: "Groq batch fixture request failed" + ├── groq-batch-collect-only-operation + │ metadata: { + │ "operation": "batch-collect-only", + │ "testRunId": "" + │ } + │ └── groq.batch [task] + │ metadata: { + │ "provider": "groq" + │ } + │ └── groq.chat.completions.create [llm] + │ input: [ + │ { + │ "content": "Collect this response.", + │ "role": "user" + │ } + │ ] + │ output: [ + │ { + │ "finish_reason": "stop", + │ "index": 0, + │ "message": { + │ "content": "COLLECTED", + │ "role": "assistant" + │ } + │ } + │ ] + │ metadata: { + │ "model": "llama-3.3-70b-versatile", + │ "provider": "groq", + │ "temperature": 0 + │ } + │ metrics: { + │ "completion_tokens": 1, + │ "prompt_tokens": 4, + │ "tokens": 5 + │ } + └── groq-batch-submission-failure-operation metadata: { - "operation": "tool", + "operation": "batch-submission-failure", "testRunId": "" } - └── groq.chat.completions.create [llm] - input: [ - { - "content": "Check the weather in Vienna and use the weather tool.", - "role": "user" - } - ] - output: [ - { - "finish_reason": "tool_calls", - "index": 0, - "logprobs": null, - "message": { - "role": "assistant", - "tool_calls": [ - { - "function": { - "arguments": "{\"location\":\"Vienna\"}", - "name": "get_weather" - }, - "id": "", - "type": "function" - } - ] - } - } - ] + └── groq.batch [task] metadata: { - "model": "llama-3.3-70b-versatile", - "provider": "groq", - "temperature": 0, - "tool_choice": { - "function": { - "name": "get_weather" - }, - "type": "function" - }, - "tools": [ - { - "function": { - "description": "Get the weather for a city.", - "name": "get_weather", - "parameters": { - "properties": { - "location": { - "description": "City name.", - "type": "string" - } - }, - "required": [ - "location" - ], - "type": "object" - } - }, - "type": "function" - } - ] - } - metrics: { - "completion_time": 0.035807806, - "completion_tokens": 11, - "prompt_time": 0.025119244, - "prompt_tokens": 246, - "queue_time": 0.145082373, - "time_to_first_token": 0, - "tokens": 257, - "total_time": 0.06092705 + "provider": "groq" } + error: "Groq batch submission failed" + ├── groq.chat.completions.create [llm] + │ input: [ + │ { + │ "content": "First pending request.", + │ "role": "user" + │ } + │ ] + │ metadata: { + │ "model": "llama-3.3-70b-versatile", + │ "provider": "groq", + │ "temperature": 0 + │ } + │ error: "Groq batch submission failed" + └── groq.chat.completions.create [llm] + input: [ + { + "content": "Second pending request.", + "role": "user" + } + ] + metadata: { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + } + error: "Groq batch submission failed" diff --git a/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-latest-wrapped.span-tree.json b/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-latest-wrapped.span-tree.json index 01ec3abcd..3affda15a 100644 --- a/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-latest-wrapped.span-tree.json +++ b/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-latest-wrapped.span-tree.json @@ -35,14 +35,10 @@ "temperature": 0 }, "metrics": { - "completion_time": 0.007446148, "completion_tokens": 2, - "prompt_time": 0.001747341, "prompt_tokens": 40, - "queue_time": 0.036266065, "time_to_first_token": 0, - "tokens": 42, - "total_time": 0.009193489 + "tokens": 42 } } ], @@ -78,18 +74,13 @@ "metadata": { "model": "llama-3.3-70b-versatile", "provider": "groq", - "stream": true, "temperature": 0 }, "metrics": { - "completion_time": 0.008458721, "completion_tokens": 2, - "prompt_time": 0.001661312, "prompt_tokens": 40, - "queue_time": 0.036348217, "time_to_first_token": 0, - "tokens": 42, - "total_time": 0.010120033 + "tokens": 42 } } ], @@ -128,19 +119,14 @@ "model": "openai/gpt-oss-20b", "provider": "groq", "reasoning_format": "parsed", - "stream": true, "temperature": 0.6 }, "metrics": { "completion_reasoning_tokens": 173, - "completion_time": 0.185192367, "completion_tokens": 183, - "prompt_time": 0.005113732, "prompt_tokens": 108, - "queue_time": 0.07155807, "time_to_first_token": 0, - "tokens": 291, - "total_time": 0.190306099 + "tokens": 291 } } ], @@ -215,14 +201,10 @@ ] }, "metrics": { - "completion_time": 0.035807806, "completion_tokens": 11, - "prompt_time": 0.025119244, "prompt_tokens": 246, - "queue_time": 0.145082373, "time_to_first_token": 0, - "tokens": 257, - "total_time": 0.06092705 + "tokens": 257 } } ], @@ -230,6 +212,207 @@ "operation": "tool", "testRunId": "" } + }, + { + "name": "groq-batch-operation", + "children": [ + { + "name": "groq.batch", + "type": "task", + "children": [ + { + "name": "groq.chat.completions.create", + "type": "llm", + "children": [], + "input": [ + { + "content": "Reply with ALPHA.", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "ALPHA", + "role": "assistant" + } + } + ], + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "metrics": { + "completion_tokens": 1, + "prompt_cached_tokens": 3, + "prompt_tokens": 5, + "tokens": 6 + } + }, + { + "name": "groq.chat.completions.create", + "type": "llm", + "children": [], + "input": [ + { + "content": "Reply with BRAVO.", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "BRAVO", + "role": "assistant" + } + } + ], + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "metrics": { + "completion_tokens": 1, + "prompt_cached_tokens": 3, + "prompt_tokens": 5, + "tokens": 6 + } + }, + { + "name": "groq.chat.completions.create", + "type": "llm", + "children": [], + "input": [ + { + "content": "This request should fail.", + "role": "user" + } + ], + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "error": "Groq batch fixture request failed" + } + ], + "metadata": { + "provider": "groq" + } + } + ], + "metadata": { + "operation": "batch", + "testRunId": "" + } + }, + { + "name": "groq-batch-collect-only-operation", + "children": [ + { + "name": "groq.batch", + "type": "task", + "children": [ + { + "name": "groq.chat.completions.create", + "type": "llm", + "children": [], + "input": [ + { + "content": "Collect this response.", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "COLLECTED", + "role": "assistant" + } + } + ], + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "metrics": { + "completion_tokens": 1, + "prompt_tokens": 4, + "tokens": 5 + } + } + ], + "metadata": { + "provider": "groq" + } + } + ], + "metadata": { + "operation": "batch-collect-only", + "testRunId": "" + } + }, + { + "name": "groq-batch-submission-failure-operation", + "children": [ + { + "name": "groq.batch", + "type": "task", + "children": [ + { + "name": "groq.chat.completions.create", + "type": "llm", + "children": [], + "input": [ + { + "content": "First pending request.", + "role": "user" + } + ], + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "error": "Groq batch submission failed" + }, + { + "name": "groq.chat.completions.create", + "type": "llm", + "children": [], + "input": [ + { + "content": "Second pending request.", + "role": "user" + } + ], + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "error": "Groq batch submission failed" + } + ], + "metadata": { + "provider": "groq" + }, + "error": "Groq batch submission failed" + } + ], + "metadata": { + "operation": "batch-submission-failure", + "testRunId": "" + } } ], "metadata": { diff --git a/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-latest-wrapped.span-tree.txt b/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-latest-wrapped.span-tree.txt index 215dd4248..3a3ff791e 100644 --- a/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-latest-wrapped.span-tree.txt +++ b/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-latest-wrapped.span-tree.txt @@ -34,14 +34,10 @@ span_tree: │ "temperature": 0 │ } │ metrics: { - │ "completion_time": 0.007446148, │ "completion_tokens": 2, - │ "prompt_time": 0.001747341, │ "prompt_tokens": 40, - │ "queue_time": 0.036266065, │ "time_to_first_token": 0, - │ "tokens": 42, - │ "total_time": 0.009193489 + │ "tokens": 42 │ } ├── groq-stream-operation │ metadata: { @@ -69,18 +65,13 @@ span_tree: │ metadata: { │ "model": "llama-3.3-70b-versatile", │ "provider": "groq", - │ "stream": true, │ "temperature": 0 │ } │ metrics: { - │ "completion_time": 0.008458721, │ "completion_tokens": 2, - │ "prompt_time": 0.001661312, │ "prompt_tokens": 40, - │ "queue_time": 0.036348217, │ "time_to_first_token": 0, - │ "tokens": 42, - │ "total_time": 0.010120033 + │ "tokens": 42 │ } ├── groq-reasoning-stream-operation │ metadata: { @@ -111,91 +102,232 @@ span_tree: │ "model": "openai/gpt-oss-20b", │ "provider": "groq", │ "reasoning_format": "parsed", - │ "stream": true, │ "temperature": 0.6 │ } │ metrics: { │ "completion_reasoning_tokens": 173, - │ "completion_time": 0.185192367, │ "completion_tokens": 183, - │ "prompt_time": 0.005113732, │ "prompt_tokens": 108, - │ "queue_time": 0.07155807, │ "time_to_first_token": 0, - │ "tokens": 291, - │ "total_time": 0.190306099 + │ "tokens": 291 │ } - └── groq-tool-operation + ├── groq-tool-operation + │ metadata: { + │ "operation": "tool", + │ "testRunId": "" + │ } + │ └── groq.chat.completions.create [llm] + │ input: [ + │ { + │ "content": "Check the weather in Vienna and use the weather tool.", + │ "role": "user" + │ } + │ ] + │ output: [ + │ { + │ "finish_reason": "tool_calls", + │ "index": 0, + │ "logprobs": null, + │ "message": { + │ "role": "assistant", + │ "tool_calls": [ + │ { + │ "function": { + │ "arguments": "{\"location\":\"Vienna\"}", + │ "name": "get_weather" + │ }, + │ "id": "", + │ "type": "function" + │ } + │ ] + │ } + │ } + │ ] + │ metadata: { + │ "model": "llama-3.3-70b-versatile", + │ "provider": "groq", + │ "temperature": 0, + │ "tool_choice": { + │ "function": { + │ "name": "get_weather" + │ }, + │ "type": "function" + │ }, + │ "tools": [ + │ { + │ "function": { + │ "description": "Get the weather for a city.", + │ "name": "get_weather", + │ "parameters": { + │ "properties": { + │ "location": { + │ "description": "City name.", + │ "type": "string" + │ } + │ }, + │ "required": [ + │ "location" + │ ], + │ "type": "object" + │ } + │ }, + │ "type": "function" + │ } + │ ] + │ } + │ metrics: { + │ "completion_tokens": 11, + │ "prompt_tokens": 246, + │ "time_to_first_token": 0, + │ "tokens": 257 + │ } + ├── groq-batch-operation + │ metadata: { + │ "operation": "batch", + │ "testRunId": "" + │ } + │ └── groq.batch [task] + │ metadata: { + │ "provider": "groq" + │ } + │ ├── groq.chat.completions.create [llm] + │ │ input: [ + │ │ { + │ │ "content": "Reply with ALPHA.", + │ │ "role": "user" + │ │ } + │ │ ] + │ │ output: [ + │ │ { + │ │ "finish_reason": "stop", + │ │ "index": 0, + │ │ "message": { + │ │ "content": "ALPHA", + │ │ "role": "assistant" + │ │ } + │ │ } + │ │ ] + │ │ metadata: { + │ │ "model": "llama-3.3-70b-versatile", + │ │ "provider": "groq", + │ │ "temperature": 0 + │ │ } + │ │ metrics: { + │ │ "completion_tokens": 1, + │ │ "prompt_cached_tokens": 3, + │ │ "prompt_tokens": 5, + │ │ "tokens": 6 + │ │ } + │ ├── groq.chat.completions.create [llm] + │ │ input: [ + │ │ { + │ │ "content": "Reply with BRAVO.", + │ │ "role": "user" + │ │ } + │ │ ] + │ │ output: [ + │ │ { + │ │ "finish_reason": "stop", + │ │ "index": 0, + │ │ "message": { + │ │ "content": "BRAVO", + │ │ "role": "assistant" + │ │ } + │ │ } + │ │ ] + │ │ metadata: { + │ │ "model": "llama-3.3-70b-versatile", + │ │ "provider": "groq", + │ │ "temperature": 0 + │ │ } + │ │ metrics: { + │ │ "completion_tokens": 1, + │ │ "prompt_cached_tokens": 3, + │ │ "prompt_tokens": 5, + │ │ "tokens": 6 + │ │ } + │ └── groq.chat.completions.create [llm] + │ input: [ + │ { + │ "content": "This request should fail.", + │ "role": "user" + │ } + │ ] + │ metadata: { + │ "model": "llama-3.3-70b-versatile", + │ "provider": "groq", + │ "temperature": 0 + │ } + │ error: "Groq batch fixture request failed" + ├── groq-batch-collect-only-operation + │ metadata: { + │ "operation": "batch-collect-only", + │ "testRunId": "" + │ } + │ └── groq.batch [task] + │ metadata: { + │ "provider": "groq" + │ } + │ └── groq.chat.completions.create [llm] + │ input: [ + │ { + │ "content": "Collect this response.", + │ "role": "user" + │ } + │ ] + │ output: [ + │ { + │ "finish_reason": "stop", + │ "index": 0, + │ "message": { + │ "content": "COLLECTED", + │ "role": "assistant" + │ } + │ } + │ ] + │ metadata: { + │ "model": "llama-3.3-70b-versatile", + │ "provider": "groq", + │ "temperature": 0 + │ } + │ metrics: { + │ "completion_tokens": 1, + │ "prompt_tokens": 4, + │ "tokens": 5 + │ } + └── groq-batch-submission-failure-operation metadata: { - "operation": "tool", + "operation": "batch-submission-failure", "testRunId": "" } - └── groq.chat.completions.create [llm] - input: [ - { - "content": "Check the weather in Vienna and use the weather tool.", - "role": "user" - } - ] - output: [ - { - "finish_reason": "tool_calls", - "index": 0, - "logprobs": null, - "message": { - "role": "assistant", - "tool_calls": [ - { - "function": { - "arguments": "{\"location\":\"Vienna\"}", - "name": "get_weather" - }, - "id": "", - "type": "function" - } - ] - } - } - ] + └── groq.batch [task] metadata: { - "model": "llama-3.3-70b-versatile", - "provider": "groq", - "temperature": 0, - "tool_choice": { - "function": { - "name": "get_weather" - }, - "type": "function" - }, - "tools": [ - { - "function": { - "description": "Get the weather for a city.", - "name": "get_weather", - "parameters": { - "properties": { - "location": { - "description": "City name.", - "type": "string" - } - }, - "required": [ - "location" - ], - "type": "object" - } - }, - "type": "function" - } - ] - } - metrics: { - "completion_time": 0.035807806, - "completion_tokens": 11, - "prompt_time": 0.025119244, - "prompt_tokens": 246, - "queue_time": 0.145082373, - "time_to_first_token": 0, - "tokens": 257, - "total_time": 0.06092705 + "provider": "groq" } + error: "Groq batch submission failed" + ├── groq.chat.completions.create [llm] + │ input: [ + │ { + │ "content": "First pending request.", + │ "role": "user" + │ } + │ ] + │ metadata: { + │ "model": "llama-3.3-70b-versatile", + │ "provider": "groq", + │ "temperature": 0 + │ } + │ error: "Groq batch submission failed" + └── groq.chat.completions.create [llm] + input: [ + { + "content": "Second pending request.", + "role": "user" + } + ] + metadata: { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + } + error: "Groq batch submission failed" diff --git a/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-wrapped.span-tree.json b/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-wrapped.span-tree.json index 84ba50ea9..8b29ec775 100644 --- a/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-wrapped.span-tree.json +++ b/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-wrapped.span-tree.json @@ -35,14 +35,10 @@ "temperature": 0 }, "metrics": { - "completion_time": 0.0072439, "completion_tokens": 2, - "prompt_time": 0.001728151, "prompt_tokens": 40, - "queue_time": 0.054452998, "time_to_first_token": 0, - "tokens": 42, - "total_time": 0.008972051 + "tokens": 42 } } ], @@ -78,18 +74,13 @@ "metadata": { "model": "llama-3.3-70b-versatile", "provider": "groq", - "stream": true, "temperature": 0 }, "metrics": { - "completion_time": 0.006718869, "completion_tokens": 2, - "prompt_time": 0.003972294, "prompt_tokens": 40, - "queue_time": 0.095351581, "time_to_first_token": 0, - "tokens": 42, - "total_time": 0.010691163 + "tokens": 42 } } ], @@ -128,19 +119,14 @@ "model": "openai/gpt-oss-20b", "provider": "groq", "reasoning_format": "parsed", - "stream": true, "temperature": 0.6 }, "metrics": { "completion_reasoning_tokens": 158, - "completion_time": 0.270633239, "completion_tokens": 168, - "prompt_time": 0.005148062, "prompt_tokens": 108, - "queue_time": 0.103912089, "time_to_first_token": 0, - "tokens": 276, - "total_time": 0.275781301 + "tokens": 276 } } ], @@ -215,14 +201,10 @@ ] }, "metrics": { - "completion_time": 0.039504246, "completion_tokens": 11, - "prompt_time": 0.020812356, "prompt_tokens": 246, - "queue_time": 0.091778475, "time_to_first_token": 0, - "tokens": 257, - "total_time": 0.060316602 + "tokens": 257 } } ], @@ -230,6 +212,207 @@ "operation": "tool", "testRunId": "" } + }, + { + "name": "groq-batch-operation", + "children": [ + { + "name": "groq.batch", + "type": "task", + "children": [ + { + "name": "groq.chat.completions.create", + "type": "llm", + "children": [], + "input": [ + { + "content": "Reply with ALPHA.", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "ALPHA", + "role": "assistant" + } + } + ], + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "metrics": { + "completion_tokens": 1, + "prompt_cached_tokens": 3, + "prompt_tokens": 5, + "tokens": 6 + } + }, + { + "name": "groq.chat.completions.create", + "type": "llm", + "children": [], + "input": [ + { + "content": "Reply with BRAVO.", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "BRAVO", + "role": "assistant" + } + } + ], + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "metrics": { + "completion_tokens": 1, + "prompt_cached_tokens": 3, + "prompt_tokens": 5, + "tokens": 6 + } + }, + { + "name": "groq.chat.completions.create", + "type": "llm", + "children": [], + "input": [ + { + "content": "This request should fail.", + "role": "user" + } + ], + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "error": "Groq batch fixture request failed" + } + ], + "metadata": { + "provider": "groq" + } + } + ], + "metadata": { + "operation": "batch", + "testRunId": "" + } + }, + { + "name": "groq-batch-collect-only-operation", + "children": [ + { + "name": "groq.batch", + "type": "task", + "children": [ + { + "name": "groq.chat.completions.create", + "type": "llm", + "children": [], + "input": [ + { + "content": "Collect this response.", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "COLLECTED", + "role": "assistant" + } + } + ], + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "metrics": { + "completion_tokens": 1, + "prompt_tokens": 4, + "tokens": 5 + } + } + ], + "metadata": { + "provider": "groq" + } + } + ], + "metadata": { + "operation": "batch-collect-only", + "testRunId": "" + } + }, + { + "name": "groq-batch-submission-failure-operation", + "children": [ + { + "name": "groq.batch", + "type": "task", + "children": [ + { + "name": "groq.chat.completions.create", + "type": "llm", + "children": [], + "input": [ + { + "content": "First pending request.", + "role": "user" + } + ], + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "error": "Groq batch submission failed" + }, + { + "name": "groq.chat.completions.create", + "type": "llm", + "children": [], + "input": [ + { + "content": "Second pending request.", + "role": "user" + } + ], + "metadata": { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + }, + "error": "Groq batch submission failed" + } + ], + "metadata": { + "provider": "groq" + }, + "error": "Groq batch submission failed" + } + ], + "metadata": { + "operation": "batch-submission-failure", + "testRunId": "" + } } ], "metadata": { diff --git a/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-wrapped.span-tree.txt b/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-wrapped.span-tree.txt index 8bab80353..dfc41ece9 100644 --- a/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-wrapped.span-tree.txt +++ b/e2e/scenarios/groq-instrumentation/__snapshots__/groq-v1-wrapped.span-tree.txt @@ -34,14 +34,10 @@ span_tree: │ "temperature": 0 │ } │ metrics: { - │ "completion_time": 0.0072439, │ "completion_tokens": 2, - │ "prompt_time": 0.001728151, │ "prompt_tokens": 40, - │ "queue_time": 0.054452998, │ "time_to_first_token": 0, - │ "tokens": 42, - │ "total_time": 0.008972051 + │ "tokens": 42 │ } ├── groq-stream-operation │ metadata: { @@ -69,18 +65,13 @@ span_tree: │ metadata: { │ "model": "llama-3.3-70b-versatile", │ "provider": "groq", - │ "stream": true, │ "temperature": 0 │ } │ metrics: { - │ "completion_time": 0.006718869, │ "completion_tokens": 2, - │ "prompt_time": 0.003972294, │ "prompt_tokens": 40, - │ "queue_time": 0.095351581, │ "time_to_first_token": 0, - │ "tokens": 42, - │ "total_time": 0.010691163 + │ "tokens": 42 │ } ├── groq-reasoning-stream-operation │ metadata: { @@ -111,91 +102,232 @@ span_tree: │ "model": "openai/gpt-oss-20b", │ "provider": "groq", │ "reasoning_format": "parsed", - │ "stream": true, │ "temperature": 0.6 │ } │ metrics: { │ "completion_reasoning_tokens": 158, - │ "completion_time": 0.270633239, │ "completion_tokens": 168, - │ "prompt_time": 0.005148062, │ "prompt_tokens": 108, - │ "queue_time": 0.103912089, │ "time_to_first_token": 0, - │ "tokens": 276, - │ "total_time": 0.275781301 + │ "tokens": 276 │ } - └── groq-tool-operation + ├── groq-tool-operation + │ metadata: { + │ "operation": "tool", + │ "testRunId": "" + │ } + │ └── groq.chat.completions.create [llm] + │ input: [ + │ { + │ "content": "Check the weather in Vienna and use the weather tool.", + │ "role": "user" + │ } + │ ] + │ output: [ + │ { + │ "finish_reason": "tool_calls", + │ "index": 0, + │ "logprobs": null, + │ "message": { + │ "role": "assistant", + │ "tool_calls": [ + │ { + │ "function": { + │ "arguments": "{\"location\":\"Vienna\"}", + │ "name": "get_weather" + │ }, + │ "id": "", + │ "type": "function" + │ } + │ ] + │ } + │ } + │ ] + │ metadata: { + │ "model": "llama-3.3-70b-versatile", + │ "provider": "groq", + │ "temperature": 0, + │ "tool_choice": { + │ "function": { + │ "name": "get_weather" + │ }, + │ "type": "function" + │ }, + │ "tools": [ + │ { + │ "function": { + │ "description": "Get the weather for a city.", + │ "name": "get_weather", + │ "parameters": { + │ "properties": { + │ "location": { + │ "description": "City name.", + │ "type": "string" + │ } + │ }, + │ "required": [ + │ "location" + │ ], + │ "type": "object" + │ } + │ }, + │ "type": "function" + │ } + │ ] + │ } + │ metrics: { + │ "completion_tokens": 11, + │ "prompt_tokens": 246, + │ "time_to_first_token": 0, + │ "tokens": 257 + │ } + ├── groq-batch-operation + │ metadata: { + │ "operation": "batch", + │ "testRunId": "" + │ } + │ └── groq.batch [task] + │ metadata: { + │ "provider": "groq" + │ } + │ ├── groq.chat.completions.create [llm] + │ │ input: [ + │ │ { + │ │ "content": "Reply with ALPHA.", + │ │ "role": "user" + │ │ } + │ │ ] + │ │ output: [ + │ │ { + │ │ "finish_reason": "stop", + │ │ "index": 0, + │ │ "message": { + │ │ "content": "ALPHA", + │ │ "role": "assistant" + │ │ } + │ │ } + │ │ ] + │ │ metadata: { + │ │ "model": "llama-3.3-70b-versatile", + │ │ "provider": "groq", + │ │ "temperature": 0 + │ │ } + │ │ metrics: { + │ │ "completion_tokens": 1, + │ │ "prompt_cached_tokens": 3, + │ │ "prompt_tokens": 5, + │ │ "tokens": 6 + │ │ } + │ ├── groq.chat.completions.create [llm] + │ │ input: [ + │ │ { + │ │ "content": "Reply with BRAVO.", + │ │ "role": "user" + │ │ } + │ │ ] + │ │ output: [ + │ │ { + │ │ "finish_reason": "stop", + │ │ "index": 0, + │ │ "message": { + │ │ "content": "BRAVO", + │ │ "role": "assistant" + │ │ } + │ │ } + │ │ ] + │ │ metadata: { + │ │ "model": "llama-3.3-70b-versatile", + │ │ "provider": "groq", + │ │ "temperature": 0 + │ │ } + │ │ metrics: { + │ │ "completion_tokens": 1, + │ │ "prompt_cached_tokens": 3, + │ │ "prompt_tokens": 5, + │ │ "tokens": 6 + │ │ } + │ └── groq.chat.completions.create [llm] + │ input: [ + │ { + │ "content": "This request should fail.", + │ "role": "user" + │ } + │ ] + │ metadata: { + │ "model": "llama-3.3-70b-versatile", + │ "provider": "groq", + │ "temperature": 0 + │ } + │ error: "Groq batch fixture request failed" + ├── groq-batch-collect-only-operation + │ metadata: { + │ "operation": "batch-collect-only", + │ "testRunId": "" + │ } + │ └── groq.batch [task] + │ metadata: { + │ "provider": "groq" + │ } + │ └── groq.chat.completions.create [llm] + │ input: [ + │ { + │ "content": "Collect this response.", + │ "role": "user" + │ } + │ ] + │ output: [ + │ { + │ "finish_reason": "stop", + │ "index": 0, + │ "message": { + │ "content": "COLLECTED", + │ "role": "assistant" + │ } + │ } + │ ] + │ metadata: { + │ "model": "llama-3.3-70b-versatile", + │ "provider": "groq", + │ "temperature": 0 + │ } + │ metrics: { + │ "completion_tokens": 1, + │ "prompt_tokens": 4, + │ "tokens": 5 + │ } + └── groq-batch-submission-failure-operation metadata: { - "operation": "tool", + "operation": "batch-submission-failure", "testRunId": "" } - └── groq.chat.completions.create [llm] - input: [ - { - "content": "Check the weather in Vienna and use the weather tool.", - "role": "user" - } - ] - output: [ - { - "finish_reason": "tool_calls", - "index": 0, - "logprobs": null, - "message": { - "role": "assistant", - "tool_calls": [ - { - "function": { - "arguments": "{\"location\":\"Vienna\"}", - "name": "get_weather" - }, - "id": "", - "type": "function" - } - ] - } - } - ] + └── groq.batch [task] metadata: { - "model": "llama-3.3-70b-versatile", - "provider": "groq", - "temperature": 0, - "tool_choice": { - "function": { - "name": "get_weather" - }, - "type": "function" - }, - "tools": [ - { - "function": { - "description": "Get the weather for a city.", - "name": "get_weather", - "parameters": { - "properties": { - "location": { - "description": "City name.", - "type": "string" - } - }, - "required": [ - "location" - ], - "type": "object" - } - }, - "type": "function" - } - ] - } - metrics: { - "completion_time": 0.039504246, - "completion_tokens": 11, - "prompt_time": 0.020812356, - "prompt_tokens": 246, - "queue_time": 0.091778475, - "time_to_first_token": 0, - "tokens": 257, - "total_time": 0.060316602 + "provider": "groq" } + error: "Groq batch submission failed" + ├── groq.chat.completions.create [llm] + │ input: [ + │ { + │ "content": "First pending request.", + │ "role": "user" + │ } + │ ] + │ metadata: { + │ "model": "llama-3.3-70b-versatile", + │ "provider": "groq", + │ "temperature": 0 + │ } + │ error: "Groq batch submission failed" + └── groq.chat.completions.create [llm] + input: [ + { + "content": "Second pending request.", + "role": "user" + } + ] + metadata: { + "model": "llama-3.3-70b-versatile", + "provider": "groq", + "temperature": 0 + } + error: "Groq batch submission failed" diff --git a/e2e/scenarios/groq-instrumentation/assertions.ts b/e2e/scenarios/groq-instrumentation/assertions.ts index 45654909c..b49158e00 100644 --- a/e2e/scenarios/groq-instrumentation/assertions.ts +++ b/e2e/scenarios/groq-instrumentation/assertions.ts @@ -34,6 +34,22 @@ function findGroqSpan( return spans.find((candidate) => candidate.output !== undefined) ?? spans[0]; } +function batchPrompt(span: CapturedLogEvent): string { + const firstMessage = Array.isArray(span.input) ? span.input[0] : undefined; + return typeof firstMessage?.content === "string" ? firstMessage.content : ""; +} + +function findBatchTrace(events: CapturedLogEvent[], operationName: string) { + const operation = findLatestSpan(events, operationName); + const task = findChildSpans(events, "groq.batch", operation?.span.id)[0]; + const children = findChildSpans( + events, + "groq.chat.completions.create", + task?.span.id, + ).sort((left, right) => batchPrompt(left).localeCompare(batchPrompt(right))); + return { operation, task, children }; +} + function spanTreeEvents(events: CapturedLogEvent[]): CapturedLogEvent[] { const chatOperation = findLatestSpan(events, "groq-chat-operation"); const streamOperation = findLatestSpan(events, "groq-stream-operation"); @@ -42,6 +58,15 @@ function spanTreeEvents(events: CapturedLogEvent[]): CapturedLogEvent[] { "groq-reasoning-stream-operation", ); const toolOperation = findLatestSpan(events, "groq-tool-operation"); + const batchTrace = findBatchTrace(events, "groq-batch-operation"); + const collectOnlyBatchTrace = findBatchTrace( + events, + "groq-batch-collect-only-operation", + ); + const failedBatchTrace = findBatchTrace( + events, + "groq-batch-submission-failure-operation", + ); return [ findLatestSpan(events, ROOT_NAME), @@ -69,6 +94,15 @@ function spanTreeEvents(events: CapturedLogEvent[]): CapturedLogEvent[] { toolOperation?.span.id, "groq.chat.completions.create", ), + batchTrace.operation, + batchTrace.task, + ...batchTrace.children, + collectOnlyBatchTrace.operation, + collectOnlyBatchTrace.task, + ...collectOnlyBatchTrace.children, + failedBatchTrace.operation, + failedBatchTrace.task, + ...failedBatchTrace.children, ].map((event) => event!); } @@ -186,6 +220,95 @@ export function defineGroqInstrumentationAssertions(options: { ); }); + test( + "captures resumable Groq Batch task and child spans", + testConfig, + () => { + const { operation, task, children } = findBatchTrace( + events, + "groq-batch-operation", + ); + + expect(task?.span.parentIds).toEqual([operation?.span.id ?? ""]); + expect(task?.row.span_attributes?.type).toBe("task"); + expect(task?.row.metadata).toMatchObject({ provider: "groq" }); + expect(task?.input).toBeUndefined(); + expect(task?.output).toBeUndefined(); + expect(children).toHaveLength(3); + expect(children.map(batchPrompt)).toEqual([ + "Reply with ALPHA.", + "Reply with BRAVO.", + "This request should fail.", + ]); + + for (const child of children) { + expect(child.span.parentIds).toEqual([task?.span.id ?? ""]); + expect(child.row.metadata).toMatchObject({ + model: expect.any(String), + provider: "groq", + }); + expect(child.metrics?.time_to_first_token).toBeUndefined(); + } + expect(children[0]?.output?.[0]?.message?.content).toBe("ALPHA"); + expect(children[1]?.output?.[0]?.message?.content).toBe("BRAVO"); + expect(children[0]?.metrics).toMatchObject({ + completion_tokens: 1, + prompt_cached_tokens: 3, + prompt_tokens: 5, + tokens: 6, + }); + expect(children[2]?.row.error).toContain( + "Groq batch fixture request failed", + ); + }, + ); + + test( + "closes pending spans after batch submission failure", + testConfig, + () => { + const { operation, task, children } = findBatchTrace( + events, + "groq-batch-submission-failure-operation", + ); + + expect(task?.span.parentIds).toEqual([operation?.span.id ?? ""]); + expect(task?.row.error).toContain("Groq batch submission failed"); + expect(children).toHaveLength(2); + expect(children.map(batchPrompt)).toEqual([ + "First pending request.", + "Second pending request.", + ]); + for (const child of children) { + expect(child.span.parentIds).toEqual([task?.span.id ?? ""]); + expect(child.row.error).toContain("Groq batch submission failed"); + expect(child.metrics?.end).toEqual(expect.any(Number)); + } + }, + ); + + test("captures collect-only Groq Batch spans", testConfig, () => { + const { operation, task, children } = findBatchTrace( + events, + "groq-batch-collect-only-operation", + ); + + expect(task?.span.parentIds).toEqual([operation?.span.id ?? ""]); + expect(operation?.span.rootId).toEqual(expect.any(String)); + expect(task?.span.rootId).toBe(operation?.span.rootId); + expect(task?.metrics?.end).toEqual(expect.any(Number)); + expect(children).toHaveLength(1); + expect(children[0]?.input).toEqual([ + { role: "user", content: "Collect this response." }, + ]); + expect(children[0]?.output?.[0]?.message?.content).toBe("COLLECTED"); + expect(children[0]?.metrics).toMatchObject({ + completion_tokens: 1, + prompt_tokens: 4, + tokens: 5, + }); + }); + test("matches span tree snapshot", testConfig, async () => { await matchSpanTreeSnapshot(events, spanSnapshotPath); }); diff --git a/e2e/scenarios/groq-instrumentation/scenario.impl.mjs b/e2e/scenarios/groq-instrumentation/scenario.impl.mjs index 3e9fc8cd0..d600782b0 100644 --- a/e2e/scenarios/groq-instrumentation/scenario.impl.mjs +++ b/e2e/scenarios/groq-instrumentation/scenario.impl.mjs @@ -1,4 +1,10 @@ -import { wrapGroq } from "braintrust"; +import { + bindGroqBatchTrace, + collectGroqBatchTrace, + failGroqBatchTrace, + startGroqBatchTrace, + wrapGroq, +} from "braintrust"; import { collectAsync, runOperation, @@ -37,6 +43,54 @@ function getWeatherToolDefinition() { }; } +function createMockBatchClient(options) { + const baseClient = new options.Groq({ + apiKey: "test-groq-key", + baseURL: "https://example.test", + maxRetries: 0, + fetch: async (_url, init) => { + const params = JSON.parse(String(init?.body)); + return new Response( + JSON.stringify({ + id: "batch_groq_e2e_fixture", + object: "batch", + endpoint: params.endpoint, + input_file_id: params.input_file_id, + completion_window: params.completion_window, + status: "validating", + created_at: 1_740_000_000, + metadata: params.metadata, + request_counts: { completed: 0, failed: 0, total: 0 }, + }), + { + headers: { "content-type": "application/json" }, + status: 200, + }, + ); + }, + }); + return options.decorateClient + ? options.decorateClient(baseClient) + : baseClient; +} + +function batchInput(items) { + return items + .map((item) => + JSON.stringify({ + custom_id: item.customId, + method: "POST", + url: "/v1/chat/completions", + body: { + model: CHAT_MODEL, + messages: [{ role: "user", content: item.prompt }], + temperature: 0, + }, + }), + ) + .join("\n"); +} + export async function runGroqInstrumentationScenario(options) { const apiKey = getApiKey(); if (!apiKey) { @@ -50,6 +104,7 @@ export async function runGroqInstrumentationScenario(options) { const client = options.decorateClient ? options.decorateClient(baseClient) : baseClient; + const batchFixtureClient = createMockBatchClient(options); await runTracedScenario({ callback: async () => { @@ -113,6 +168,173 @@ export async function runGroqInstrumentationScenario(options) { tools: [getWeatherToolDefinition()], }); }); + + await runOperation("groq-batch-operation", "batch", async () => { + const items = [ + { customId: "batch_alpha", prompt: "Reply with ALPHA." }, + { customId: "batch_bravo", prompt: "Reply with BRAVO." }, + { customId: "batch_error", prompt: "This request should fail." }, + ]; + const input = batchInput(items); + const params = await startGroqBatchTrace({ + inputFile: { id: "file_groq_batch_fixture" }, + input, + params: { + completion_window: "48h", + endpoint: "/v1/chat/completions", + }, + }); + const created = await batchFixtureClient.batches.create(params); + const traceContext = await bindGroqBatchTrace({ batch: created }); + if (!traceContext) { + throw new Error("Expected the accepted Groq batch trace to bind"); + } + const completedAt = Date.now() / 1000 + 60; + const outputFile = Promise.resolve( + new Response( + [items[1], items[0]] + .map((item, index) => + JSON.stringify({ + custom_id: item.customId, + response: { + status_code: 200, + body: { + model: CHAT_MODEL, + choices: [ + { + index: 0, + finish_reason: "stop", + message: { + role: "assistant", + content: index === 0 ? "BRAVO" : "ALPHA", + }, + }, + ], + usage: { + prompt_tokens: 5, + completion_tokens: 1, + total_tokens: 6, + }, + x_groq: { + usage: { + dram_cached_tokens: 2, + sram_cached_tokens: 1, + }, + }, + }, + }, + }), + ) + .join("\n"), + ), + ); + const errorFile = Promise.resolve( + new Response( + JSON.stringify({ + custom_id: items[2].customId, + error: { message: "Groq batch fixture request failed" }, + }), + ), + ); + await collectGroqBatchTrace({ + batch: { + ...created, + status: "completed", + completed_at: completedAt, + request_counts: { completed: 2, failed: 1, total: 3 }, + }, + traceContext, + inputFile: input, + outputFile, + errorFile, + }); + if (!(await outputFile).bodyUsed || !(await errorFile).bodyUsed) { + throw new Error( + "Expected Groq batch result responses to be consumed", + ); + } + }); + + await runOperation( + "groq-batch-collect-only-operation", + "batch-collect-only", + async () => { + const items = [ + { customId: "collect_only", prompt: "Collect this response." }, + ]; + const input = batchInput(items); + const completedAt = Date.now() / 1000; + await collectGroqBatchTrace({ + batch: { + id: "batch_groq_collect_only_fixture", + endpoint: "/v1/chat/completions", + input_file_id: "file_groq_collect_only_fixture", + status: "completed", + created_at: completedAt, + completed_at: completedAt, + request_counts: { completed: 1, failed: 0, total: 1 }, + }, + inputFile: input, + outputFile: new Response( + JSON.stringify({ + custom_id: items[0].customId, + response: { + status_code: 200, + body: { + model: CHAT_MODEL, + choices: [ + { + index: 0, + finish_reason: "stop", + message: { + role: "assistant", + content: "COLLECTED", + }, + }, + ], + usage: { + prompt_tokens: 4, + completion_tokens: 1, + total_tokens: 5, + }, + }, + }, + }), + ), + }); + }, + ); + + await runOperation( + "groq-batch-submission-failure-operation", + "batch-submission-failure", + async () => { + const items = [ + { customId: "submission_one", prompt: "First pending request." }, + { customId: "submission_two", prompt: "Second pending request." }, + ]; + const input = batchInput(items); + const params = await startGroqBatchTrace({ + inputFile: { id: "file_groq_batch_submission_failure" }, + input, + params: { + completion_window: "24h", + endpoint: "/v1/chat/completions", + }, + }); + try { + await { + batches: { + create: async () => { + throw new Error("Groq batch submission failed"); + }, + }, + }.batches.create(params); + } catch (error) { + await failGroqBatchTrace({ params, input, error }); + } + }, + ); }, metadata: { scenario: SCENARIO_NAME, diff --git a/js/src/exports.ts b/js/src/exports.ts index 4072fcc54..726efb39c 100644 --- a/js/src/exports.ts +++ b/js/src/exports.ts @@ -225,6 +225,27 @@ export { wrapOllama } from "./wrappers/ollama"; export { wrapCohere } from "./wrappers/cohere"; export { wrapVoyageAI } from "./wrappers/voyageai"; export { wrapGroq } from "./wrappers/groq"; +export { + bindGroqBatchTrace, + collectGroqBatchTrace, + failGroqBatchTrace, + startGroqBatchTrace, +} from "./groq-batch"; +export type { + BindGroqBatchTraceArgs, + CollectGroqBatchTraceArgs, + FailGroqBatchTraceArgs, + GroqBatchCreateInputParams, + GroqBatchCreateParams, + GroqBatchFile, + GroqBatchFileFactory, + GroqBatchJSONL, + GroqBatchLike, + GroqBatchReplayableJSONL, + GroqBatchTraceContext, + GroqFileLike, + StartGroqBatchTraceArgs, +} from "./groq-batch-types"; export { wrapBedrockRuntime } from "./wrappers/bedrock-runtime"; export { wrapCopilotClient } from "./wrappers/github-copilot"; export { diff --git a/js/src/groq-batch-types.ts b/js/src/groq-batch-types.ts new file mode 100644 index 000000000..4e53269e2 --- /dev/null +++ b/js/src/groq-batch-types.ts @@ -0,0 +1,82 @@ +export type GroqBatchJSONL = + | string + | Iterable + | AsyncIterable; + +export type GroqBatchReplayableJSONL = + | string + | (() => Iterable | AsyncIterable); + +type GroqBatchFileContent = GroqBatchJSONL | Response; + +export type GroqBatchFileFactory = () => + | GroqBatchFileContent + | PromiseLike; + +export type GroqBatchFile = + | GroqBatchFileContent + | GroqBatchFileFactory + | PromiseLike; + +export interface GroqBatchLike { + id: string; + endpoint: string; + input_file_id: string; + status: string; + created_at?: number | null; + completed_at?: number | null; + failed_at?: number | null; + expired_at?: number | null; + cancelled_at?: number | null; + metadata?: unknown | null; + request_counts?: { + completed?: number; + failed?: number; + total?: number; + } | null; +} + +export interface GroqBatchCreateParams { + input_file_id: string; + endpoint: "/v1/chat/completions"; + completion_window: string; + metadata?: Record | null; +} + +export interface GroqBatchCreateInputParams { + endpoint: "/v1/chat/completions"; + completion_window: string; + metadata?: Record | null; +} + +export interface GroqFileLike { + id: string; +} + +export interface StartGroqBatchTraceArgs { + inputFile: GroqFileLike; + input: GroqBatchReplayableJSONL; + params: GroqBatchCreateInputParams; +} + +export interface GroqBatchTraceContext { + readonly value: string; +} + +export interface BindGroqBatchTraceArgs { + batch: TBatch; +} + +export interface CollectGroqBatchTraceArgs { + batch: TBatch; + traceContext?: GroqBatchTraceContext; + inputFile: GroqBatchFile; + outputFile?: GroqBatchFile; + errorFile?: GroqBatchFile; +} + +export interface FailGroqBatchTraceArgs { + params: GroqBatchCreateParams; + input: GroqBatchReplayableJSONL; + error: unknown; +} diff --git a/js/src/groq-batch.ts b/js/src/groq-batch.ts new file mode 100644 index 000000000..2300b907b --- /dev/null +++ b/js/src/groq-batch.ts @@ -0,0 +1,116 @@ +import { groqChannels } from "./instrumentation/plugins/groq-channels"; +import type { + BindGroqBatchTraceArgs, + CollectGroqBatchTraceArgs, + FailGroqBatchTraceArgs, + GroqBatchCreateParams, + GroqBatchLike, + GroqBatchTraceContext, + StartGroqBatchTraceArgs, +} from "./groq-batch-types"; + +/** + * Start a trace for a new Groq Batch, with pending spans for its requests, and + * return the parameters to pass to `groq.batches.create()`. + * + * Only `/v1/chat/completions` batches are supported. + * Iterable inputs must be supplied through a factory that returns a fresh + * iterable on every call so submission failures can resume the pending spans. + * Instrumentation enforces Groq's 50,000-record and 100 MB batch limits. + * Inputs are validated in one pass and replayed for bounded streaming span + * emission, supporting the full provider file limit. + */ +export async function startGroqBatchTrace( + args: StartGroqBatchTraceArgs, +): Promise { + return await groqChannels.batchesStartTrace.invoke( + async (input) => ({ + ...input.params, + input_file_id: input.inputFile.id, + }), + undefined, + [args], + {}, + ); +} + +/** + * Bind an accepted Groq Batch ID to the pending trace created by + * `startGroqBatchTrace()`. + * + * Call this with the object returned by `groq.batches.create()`, then retain + * the opaque result for `collectGroqBatchTrace()`. This helper performs no Groq + * API requests and returns `undefined` when the batch has no valid pending + * Braintrust context. + */ +export async function bindGroqBatchTrace( + args: BindGroqBatchTraceArgs, +): Promise { + const result = await groqChannels.batchesBindTrace.invoke( + async (): Promise<{ traceContext?: GroqBatchTraceContext }> => ({}), + undefined, + [args], + {}, + ); + return result.traceContext; +} + +/** + * Collect a Groq Batch trace from caller-supplied Batch API data. + * + * This helper performs no Groq API requests. Supplied Response bodies are + * consumed directly. Completed batches are validated and staged atomically in + * a bounded 16 MiB result buffer before span updates are emitted. Result-file + * factories are validated in one pass and replayed for bounded streaming + * emission, supporting the full 100 MB provider limit. Without signed start + * context, collect-only tracing also uses a bounded 16 MiB normalized-input + * buffer and otherwise skips with a debug diagnostic. Files are limited to + * Groq's 50,000-record and 100 MB limits. A trace started with + * `startGroqBatchTrace()` also requires the opaque context returned by + * `bindGroqBatchTrace()`. Local collection failures leave completed-batch + * spans pending so callers can retry with fresh or replayable file sources. + */ +export async function collectGroqBatchTrace( + args: CollectGroqBatchTraceArgs, +): Promise { + const inputFile = Promise.resolve(args.inputFile); + const outputFile = + args.outputFile === undefined + ? undefined + : Promise.resolve(args.outputFile); + const errorFile = + args.errorFile === undefined ? undefined : Promise.resolve(args.errorFile); + + void inputFile.catch(() => undefined); + void outputFile?.catch(() => undefined); + void errorFile?.catch(() => undefined); + + const batch = await groqChannels.batchesCollectTrace.invoke( + async (input) => input.batch, + undefined, + [{ ...args, inputFile, outputFile, errorFile }], + {}, + ); + // The channel preserves the exact caller-supplied batch object, including + // fields from newer Groq SDK versions. + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions + return batch as TBatch; +} + +/** + * End a pending Groq Batch trace after `groq.batches.create()` rejects. + * + * The original replayable input is required to validate and resume every child + * span. This helper performs no Groq API requests and never rethrows the + * supplied error. + */ +export async function failGroqBatchTrace( + args: FailGroqBatchTraceArgs, +): Promise { + await groqChannels.batchesFailTrace.invoke( + async () => undefined, + undefined, + [args], + {}, + ); +} diff --git a/js/src/instrumentation/plugins/groq-batch-instrumentation.test.ts b/js/src/instrumentation/plugins/groq-batch-instrumentation.test.ts new file mode 100644 index 000000000..f47ba1312 --- /dev/null +++ b/js/src/instrumentation/plugins/groq-batch-instrumentation.test.ts @@ -0,0 +1,1411 @@ +import { + afterEach, + beforeAll, + beforeEach, + describe, + expect, + it, + vi, +} from "vitest"; +import { debugLogger } from "../../debug-logger"; +import { + _exportsForTestingOnly, + _internalGetGlobalState, + initLogger, + traced, +} from "../../logger"; +import { configureNode } from "../../node/config"; +import { + bindGroqBatchTrace, + collectGroqBatchTrace, + failGroqBatchTrace, + startGroqBatchTrace, +} from "../../groq-batch"; +import { GroqPlugin } from "./groq-plugin"; + +try { + configureNode(); +} catch { + // Best-effort initialization for test environments. +} + +const CONTEXT_KEY = "braintrust.batch_context"; +const chatInput = [ + { + custom_id: "ok", + method: "POST", + url: "/v1/chat/completions", + body: { + model: "llama-test", + messages: [{ role: "user", content: "hi" }], + temperature: 0, + user: "not-allowlisted", + }, + }, + { + custom_id: "bad", + method: "POST", + url: "/v1/chat/completions", + body: { + model: "llama-test", + messages: [{ role: "user", content: "fail" }], + }, + }, +]; + +function jsonl(records: unknown[]): string { + return records.map((record) => JSON.stringify(record)).join("\n"); +} + +async function startBatch( + input: string | (() => Iterable) = jsonl(chatInput), +) { + const params = await startGroqBatchTrace({ + inputFile: { id: "file_input" }, + input, + params: { + endpoint: "/v1/chat/completions", + completion_window: "48h", + metadata: { caller: "value" }, + }, + }); + const batch = { + id: "batch_test", + ...params, + status: "validating", + }; + const traceContext = await bindGroqBatchTrace({ batch }); + if (!traceContext) { + throw new Error("Expected the Groq batch trace to bind"); + } + return { + input, + params, + batch, + traceContext, + }; +} + +function chunkedResponse(records: unknown[]): Response { + const bytes = new TextEncoder().encode(jsonl(records)); + let offset = 0; + return new Response( + new ReadableStream({ + pull(controller) { + if (offset < bytes.length) { + controller.enqueue(bytes.slice(offset, offset + 7)); + offset += 7; + } else { + controller.close(); + } + }, + }), + ); +} + +describe("Groq Batch instrumentation", () => { + let backgroundLogger: ReturnType< + typeof _exportsForTestingOnly.useTestBackgroundLogger + >; + let plugin: GroqPlugin; + + beforeAll(async () => { + await _exportsForTestingOnly.simulateLoginForTests(); + _internalGetGlobalState()._internalSetTraceContextSigningSecret( + "groq-batch-test-secret", + ); + }); + + beforeEach(() => { + backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); + initLogger({ + projectName: "groq-batch-instrumentation.test.ts", + projectId: "test-project-id", + }); + plugin = new GroqPlugin(); + plugin.enable(); + }); + + afterEach(() => { + plugin.disable(); + _exportsForTestingOnly.clearTestBackgroundLogger(); + vi.restoreAllMocks(); + }); + + it("returns provider-ready params and starts pending task and LLM spans", async () => { + const sourceMetadata = { caller: "value" }; + const params = await startGroqBatchTrace({ + inputFile: { id: "file_input" }, + input: jsonl(chatInput), + params: { + endpoint: "/v1/chat/completions", + completion_window: "48h", + metadata: sourceMetadata, + }, + }); + + expect(params).toMatchObject({ + input_file_id: "file_input", + endpoint: "/v1/chat/completions", + completion_window: "48h", + metadata: { + caller: "value", + [CONTEXT_KEY]: expect.any(String), + }, + }); + expect(sourceMetadata).toEqual({ caller: "value" }); + + const rows = (await backgroundLogger.drain()) as Array>; + const task = rows.find((row) => row.span_attributes?.name === "groq.batch"); + const children = rows.filter( + (row) => row.span_attributes?.name === "groq.chat.completions.create", + ); + expect(rows).toHaveLength(3); + expect(task).toMatchObject({ + metadata: { provider: "groq" }, + span_attributes: { type: "task" }, + span_parents: [], + }); + expect(task?.metrics).not.toHaveProperty("end"); + expect(children).toHaveLength(2); + expect( + children.every((row) => row.root_span_id === task?.root_span_id), + ).toBe(true); + expect(children.map((row) => row.input?.[0]?.content).sort()).toEqual([ + "fail", + "hi", + ]); + expect( + children.find((row) => row.input?.[0]?.content === "hi"), + ).toMatchObject({ + metadata: { + model: "llama-test", + provider: "groq", + temperature: 0, + }, + }); + expect( + children.find((row) => row.input?.[0]?.content === "hi")?.metadata, + ).not.toHaveProperty("user"); + }); + + it("flushes large batch span writes in bounded chunks", async () => { + const flushSpy = vi.spyOn(backgroundLogger, "flush"); + const input = Array.from({ length: 1001 }, (_, index) => ({ + custom_id: `request-${index}`, + method: "POST", + url: "/v1/chat/completions", + body: { + model: "llama-test", + messages: [{ role: "user", content: `prompt-${index}` }], + }, + })); + + await startGroqBatchTrace({ + inputFile: { id: "file_large" }, + input: () => input.values(), + params: { + endpoint: "/v1/chat/completions", + completion_window: "24h", + }, + }); + + expect(flushSpy.mock.calls.length).toBeGreaterThanOrEqual(4); + expect(await backgroundLogger.drain()).toHaveLength(1002); + }); + + it("preserves a collection repair path when a partial start cannot be closed", async () => { + const originalLog = backgroundLogger.log.bind(backgroundLogger); + let logCalls = 0; + vi.spyOn(backgroundLogger, "log").mockImplementation((items) => { + logCalls += 1; + if (logCalls > 1) { + throw new Error("simulated start child emission failure"); + } + originalLog(items); + }); + + const params = await startGroqBatchTrace({ + inputFile: { id: "file_incomplete_start" }, + input: jsonl(chatInput), + params: { + endpoint: "/v1/chat/completions", + completion_window: "24h", + metadata: { caller: "value" }, + }, + }); + + expect(params.metadata?.[CONTEXT_KEY]).toEqual(expect.any(String)); + expect(JSON.parse(params.metadata?.[CONTEXT_KEY] ?? "{}")).toMatchObject({ + repairInputs: true, + }); + const partialRows = (await backgroundLogger.drain()) as Array< + Record + >; + expect( + partialRows.find((row) => row.span_attributes?.name === "groq.batch") + ?.metrics?.end, + ).toBeUndefined(); + + vi.mocked(backgroundLogger.log).mockRestore(); + const batch = { + id: "batch_repaired_start", + ...params, + status: "validating", + }; + const traceContext = await bindGroqBatchTrace({ batch }); + await collectGroqBatchTrace({ + batch: { + ...batch, + status: "completed", + completed_at: Date.now() / 1000 + 10, + request_counts: { completed: 2, failed: 0, total: 2 }, + }, + traceContext, + inputFile: jsonl(chatInput), + outputFile: jsonl( + chatInput.map((record) => ({ + custom_id: record.custom_id, + response: { status_code: 200, body: { choices: [] } }, + })), + ), + }); + + const repairedRows = (await backgroundLogger.drain()) as Array< + Record + >; + expect(repairedRows).toHaveLength(3); + expect(repairedRows.every((row) => row.metrics?.end !== undefined)).toBe( + true, + ); + expect( + repairedRows + .filter( + (row) => row.span_attributes?.name === "groq.chat.completions.create", + ) + .map((row) => row.input?.[0]?.content) + .sort(), + ).toEqual(["fail", "hi"]); + }); + + it("withholds signed context when replayable start input changes", async () => { + let calls = 0; + const params = await startGroqBatchTrace({ + inputFile: { id: "file_changed_start" }, + input: () => { + calls += 1; + return calls % 2 === 1 ? chatInput : [chatInput[0]]; + }, + params: { + endpoint: "/v1/chat/completions", + completion_window: "24h", + metadata: { caller: "value" }, + }, + }); + + expect(calls).toBeGreaterThanOrEqual(2); + expect(params.metadata).toEqual({ caller: "value" }); + expect(params.metadata).not.toHaveProperty(CONTEXT_KEY); + expect(await backgroundLogger.drain()).toEqual([]); + }); + + it("completes successful and failed children from out-of-order streams", async () => { + const { batch, input, traceContext } = await startBatch(); + await backgroundLogger.drain(); + const completedAt = Date.now() / 1000 + 10; + const terminal = { + ...batch, + status: "completed", + completed_at: completedAt, + request_counts: { completed: 1, failed: 1, total: 2 }, + }; + const outputFile = chunkedResponse([ + { + custom_id: "ok", + response: { + status_code: 200, + body: { + model: "llama-resolved", + choices: [ + { + index: 0, + finish_reason: "stop", + message: { role: "assistant", content: "hello" }, + }, + ], + usage: { + prompt_tokens: 3, + completion_tokens: 1, + total_tokens: 4, + }, + x_groq: { + usage: { dram_cached_tokens: 2, sram_cached_tokens: 1 }, + }, + }, + }, + }, + ]); + const errorFile = chunkedResponse([ + { + custom_id: "bad", + error: { message: "request failed" }, + }, + ]); + + expect( + await collectGroqBatchTrace({ + batch: terminal, + traceContext, + inputFile: input, + outputFile, + errorFile, + }), + ).toBe(terminal); + expect(outputFile.bodyUsed).toBe(true); + expect(errorFile.bodyUsed).toBe(true); + + const rows = (await backgroundLogger.drain()) as Array>; + const task = rows.find((row) => row.span_attributes?.name === "groq.batch"); + const ok = rows.find( + (row) => row.output?.[0]?.message?.content === "hello", + ); + const bad = rows.find((row) => row.error?.includes("request failed")); + expect(task?.metrics?.end).toBe(completedAt); + expect(ok).toMatchObject({ + metadata: { model: "llama-resolved", provider: "groq" }, + metrics: { + completion_tokens: 1, + prompt_cached_tokens: 3, + prompt_tokens: 3, + tokens: 4, + end: completedAt, + }, + }); + expect(bad?.metrics?.end).toBe(completedAt); + }); + + it("completes a batch from an output file alone when request counts are absent", async () => { + const { batch, input, traceContext } = await startBatch(); + await backgroundLogger.drain(); + await collectGroqBatchTrace({ + batch: { + ...batch, + status: "completed", + completed_at: Date.now() / 1000 + 10, + }, + traceContext, + inputFile: input, + outputFile: jsonl( + chatInput.map((record) => ({ + custom_id: record.custom_id, + response: { status_code: 200, body: { choices: [] } }, + })), + ), + }); + + const rows = (await backgroundLogger.drain()) as Array>; + expect(rows).toHaveLength(3); + expect( + rows.find((row) => row.span_attributes?.name === "groq.batch")?.error, + ).toBeUndefined(); + expect(rows.every((row) => row.metrics?.end !== undefined)).toBe(true); + }); + + it("streams replayable completed results beyond the one-shot buffer", async () => { + const recordCount = 18; + const input = Array.from({ length: recordCount }, (_, index) => ({ + custom_id: `large-${index}`, + method: "POST", + url: "/v1/chat/completions", + body: { + model: "llama-test", + messages: [{ role: "user", content: `prompt-${index}` }], + }, + })); + const started = await startBatch(() => input.values()); + await backgroundLogger.drain(); + const largeContent = "x".repeat(1024 * 1024); + let resultFactoryCalls = 0; + const outputFile = () => { + resultFactoryCalls += 1; + return input.map((record) => ({ + custom_id: record.custom_id, + response: { + status_code: 200, + body: { + choices: [ + { message: { role: "assistant", content: largeContent } }, + ], + }, + }, + })); + }; + + await collectGroqBatchTrace({ + batch: { + ...started.batch, + status: "completed", + completed_at: Date.now() / 1000 + 10, + request_counts: { + completed: recordCount, + failed: 0, + total: recordCount, + }, + }, + traceContext: started.traceContext, + inputFile: started.input, + outputFile, + }); + + const rows = (await backgroundLogger.drain()) as Array>; + const task = rows.find((row) => row.span_attributes?.name === "groq.batch"); + expect(resultFactoryCalls).toBeGreaterThanOrEqual(2); + expect(rows).toHaveLength(recordCount + 1); + expect(task?.error).toBeUndefined(); + expect(rows.every((row) => row.metrics?.end !== undefined)).toBe(true); + }); + + it("leaves one-shot staging overflow retryable", async () => { + const recordCount = 18; + const input = Array.from({ length: recordCount }, (_, index) => ({ + custom_id: `one-shot-${index}`, + method: "POST", + url: "/v1/chat/completions", + body: { + model: "llama-test", + messages: [{ role: "user", content: `prompt-${index}` }], + }, + })); + const started = await startBatch(() => input.values()); + await backgroundLogger.drain(); + const largeContent = "x".repeat(1024 * 1024); + const outputFile = (function* () { + for (const record of input) { + yield { + custom_id: record.custom_id, + response: { + status_code: 200, + body: { + choices: [ + { message: { role: "assistant", content: largeContent } }, + ], + }, + }, + }; + } + })(); + + await collectGroqBatchTrace({ + batch: { + ...started.batch, + status: "completed", + completed_at: Date.now() / 1000 + 10, + request_counts: { + completed: recordCount, + failed: 0, + total: recordCount, + }, + }, + traceContext: started.traceContext, + inputFile: started.input, + outputFile, + }); + + expect(await backgroundLogger.drain()).toEqual([]); + + await collectGroqBatchTrace({ + batch: { + ...started.batch, + status: "completed", + completed_at: Date.now() / 1000 + 10, + request_counts: { + completed: recordCount, + failed: 0, + total: recordCount, + }, + }, + traceContext: started.traceContext, + inputFile: started.input, + outputFile: () => + input.map((record) => ({ + custom_id: record.custom_id, + response: { + status_code: 200, + body: { + choices: [ + { message: { role: "assistant", content: largeContent } }, + ], + }, + }, + })), + }); + + const rows = (await backgroundLogger.drain()) as Array>; + expect(rows).toHaveLength(recordCount + 1); + expect(rows.every((row) => row.error === undefined)).toBe(true); + expect(rows.every((row) => row.metrics?.end !== undefined)).toBe(true); + }); + + it("uses standard JSON semantics for iterable input records", async () => { + const serializedRecord = { + custom_id: "json-semantics", + method: "POST", + url: "/v1/chat/completions", + body: { + model: "llama-test", + messages: [{ role: "user", content: "normalized" }], + omitted: undefined, + array: [undefined], + }, + omitted: undefined, + toJSON() { + return { + custom_id: this.custom_id, + method: this.method, + url: this.url, + body: this.body, + omittedByToJSON: undefined, + }; + }, + }; + const params = await startGroqBatchTrace({ + inputFile: { id: "file_json_semantics" }, + input: () => [serializedRecord], + params: { + endpoint: "/v1/chat/completions", + completion_window: "24h", + }, + }); + const batch = { + id: "batch_json_semantics", + ...params, + status: "validating", + }; + const traceContext = await bindGroqBatchTrace({ batch }); + expect(traceContext).toMatchObject({ value: expect.any(String) }); + await backgroundLogger.drain(); + + await collectGroqBatchTrace({ + batch: { + ...batch, + status: "completed", + completed_at: Date.now() / 1000 + 10, + request_counts: { completed: 1, failed: 0, total: 1 }, + }, + traceContext, + inputFile: JSON.stringify(serializedRecord), + outputFile: jsonl([ + { + custom_id: "json-semantics", + response: { status_code: 200, body: { choices: [] } }, + }, + ]), + }); + + const rows = (await backgroundLogger.drain()) as Array>; + expect(rows).toHaveLength(2); + expect(rows.every((row) => row.metrics?.end !== undefined)).toBe(true); + }); + + it("rejects a trace context bound to a different provider batch ID", async () => { + const { batch, traceContext } = await startBatch(); + await backgroundLogger.drain(); + const inputFile = chunkedResponse(chatInput); + const outputFile = chunkedResponse([ + { + custom_id: "ok", + response: { status_code: 200, body: { choices: [] } }, + }, + { + custom_id: "bad", + response: { status_code: 200, body: { choices: [] } }, + }, + ]); + + await collectGroqBatchTrace({ + batch: { + ...batch, + id: "batch_other", + status: "completed", + completed_at: Date.now() / 1000 + 10, + request_counts: { completed: 2, failed: 0, total: 2 }, + }, + traceContext, + inputFile, + outputFile, + }); + + expect(inputFile.bodyUsed).toBe(false); + expect(outputFile.bodyUsed).toBe(false); + expect(await backgroundLogger.drain()).toEqual([]); + }); + + it("leaves child emission failures retryable", async () => { + const { batch, input, traceContext } = await startBatch(); + await backgroundLogger.drain(); + const originalLog = backgroundLogger.log.bind(backgroundLogger); + let logCalls = 0; + vi.spyOn(backgroundLogger, "log").mockImplementation((items) => { + logCalls += 1; + if (logCalls > 1) { + throw new Error("simulated child emission failure"); + } + originalLog(items); + }); + + const terminal = { + ...batch, + status: "completed", + completed_at: Date.now() / 1000 + 10, + request_counts: { completed: 2, failed: 0, total: 2 }, + }; + const outputFile = jsonl( + chatInput.map((record) => ({ + custom_id: record.custom_id, + response: { status_code: 200, body: { choices: [] } }, + })), + ); + await collectGroqBatchTrace({ + batch: terminal, + traceContext, + inputFile: input, + outputFile, + }); + + const partialRows = (await backgroundLogger.drain()) as Array< + Record + >; + const pendingTask = partialRows.find( + (row) => row.span_attributes?.name === "groq.batch", + ); + expect(pendingTask?.error).toBeUndefined(); + expect(pendingTask?.metrics?.end).toBeUndefined(); + + vi.mocked(backgroundLogger.log).mockRestore(); + await collectGroqBatchTrace({ + batch: terminal, + traceContext, + inputFile: input, + outputFile, + }); + + const rows = (await backgroundLogger.drain()) as Array>; + const task = rows.find((row) => row.span_attributes?.name === "groq.batch"); + expect(rows).toHaveLength(3); + expect(task?.error).toBeUndefined(); + expect(rows.every((row) => row.metrics?.end !== undefined)).toBe(true); + }); + + it("leaves failed batches retryable when a terminal child cannot be emitted", async () => { + const { batch, input, traceContext } = await startBatch(); + await backgroundLogger.drain(); + const originalLog = backgroundLogger.log.bind(backgroundLogger); + let logCalls = 0; + vi.spyOn(backgroundLogger, "log").mockImplementation((items) => { + logCalls += 1; + if (logCalls > 1) { + throw new Error("simulated terminal child emission failure"); + } + originalLog(items); + }); + + const terminal = { + ...batch, + status: "failed", + failed_at: Date.now() / 1000 + 10, + request_counts: { completed: 1, failed: 1, total: 2 }, + }; + const outputFile = jsonl([ + { + custom_id: "ok", + response: { status_code: 200, body: { choices: [] } }, + }, + ]); + const errorFile = jsonl([ + { custom_id: "bad", error: { message: "provider failure" } }, + ]); + await collectGroqBatchTrace({ + batch: terminal, + traceContext, + inputFile: input, + outputFile, + errorFile, + }); + + const partialRows = (await backgroundLogger.drain()) as Array< + Record + >; + const pendingTask = partialRows.find( + (row) => row.span_attributes?.name === "groq.batch", + ); + expect(pendingTask?.error).toBeUndefined(); + expect(pendingTask?.metrics?.end).toBeUndefined(); + + vi.mocked(backgroundLogger.log).mockRestore(); + await collectGroqBatchTrace({ + batch: terminal, + traceContext, + inputFile: input, + outputFile, + errorFile, + }); + + const rows = (await backgroundLogger.drain()) as Array>; + expect(rows).toHaveLength(3); + expect(rows.every((row) => row.metrics?.end !== undefined)).toBe(true); + expect( + rows.find((row) => row.span_attributes?.name === "groq.batch")?.error, + ).toContain("Groq Batch failed"); + }); + + it("leaves completed batches pending when results are incomplete", async () => { + const { batch, input, traceContext } = await startBatch(); + await backgroundLogger.drain(); + await collectGroqBatchTrace({ + batch: { + ...batch, + status: "completed", + completed_at: Date.now() / 1000 + 10, + request_counts: { completed: 1, failed: 0, total: 2 }, + }, + traceContext, + inputFile: input, + outputFile: jsonl([ + { + custom_id: "ok", + response: { + status_code: 200, + body: { choices: [], usage: { total_tokens: 0 } }, + }, + }, + ]), + }); + + expect(await backgroundLogger.drain()).toEqual([]); + }); + + it("does not consume completed result files until all required sources are supplied", async () => { + const { batch, traceContext } = await startBatch(); + await backgroundLogger.drain(); + const terminal = { + ...batch, + status: "completed", + completed_at: Date.now() / 1000 + 10, + request_counts: { completed: 1, failed: 1, total: 2 }, + }; + const outputFile = chunkedResponse([ + { + custom_id: "ok", + response: { status_code: 200, body: { choices: [] } }, + }, + ]); + const errorFile = chunkedResponse([ + { custom_id: "bad", error: { message: "atomic failure" } }, + ]); + const inputFile = chunkedResponse(chatInput); + + await collectGroqBatchTrace({ + batch: terminal, + traceContext, + inputFile, + outputFile, + }); + expect(inputFile.bodyUsed).toBe(false); + expect(outputFile.bodyUsed).toBe(false); + expect(await backgroundLogger.drain()).toEqual([]); + + await collectGroqBatchTrace({ + batch: terminal, + traceContext, + inputFile, + errorFile, + }); + expect(inputFile.bodyUsed).toBe(false); + expect(errorFile.bodyUsed).toBe(false); + expect(await backgroundLogger.drain()).toEqual([]); + + await collectGroqBatchTrace({ + batch: terminal, + traceContext, + inputFile, + outputFile, + errorFile, + }); + expect(inputFile.bodyUsed).toBe(true); + expect(outputFile.bodyUsed).toBe(true); + expect(errorFile.bodyUsed).toBe(true); + expect(await backgroundLogger.drain()).toHaveLength(3); + }); + + it("collects a complete batch under the active parent without signed context", async () => { + const diagnostic = vi + .spyOn(debugLogger, "debug") + .mockImplementation(() => {}); + const completedAt = Date.now() / 1000 + 10; + await traced( + async () => { + await collectGroqBatchTrace({ + batch: { + id: "batch_collect_only", + endpoint: "/v1/chat/completions", + input_file_id: "file_collect_only", + status: "completed", + created_at: completedAt - 10, + completed_at: completedAt, + request_counts: { completed: 1, failed: 1, total: 2 }, + }, + inputFile: jsonl(chatInput), + outputFile: jsonl([ + { + custom_id: "ok", + response: { + status_code: 200, + body: { + model: "llama-collected", + choices: [ + { + index: 0, + finish_reason: "stop", + message: { role: "assistant", content: "collected" }, + }, + ], + }, + }, + }, + ]), + errorFile: jsonl([ + { custom_id: "bad", error: { message: "collected failure" } }, + ]), + }); + }, + { name: "collect-parent" }, + ); + + const rows = (await backgroundLogger.drain()) as Array>; + const parent = rows.find( + (row) => row.span_attributes?.name === "collect-parent", + ); + const task = rows.find((row) => row.span_attributes?.name === "groq.batch"); + expect(task?.span_parents).toEqual([parent?.span_id]); + expect(parent?.root_span_id).toEqual(expect.any(String)); + expect(task?.root_span_id).toBe(parent?.root_span_id); + expect(task?.metrics?.start).toBe(completedAt - 10); + expect(task?.metrics?.end).toBe(completedAt); + expect(rows).toHaveLength(4); + expect( + rows.find((row) => row.output?.[0]?.message?.content === "collected"), + ).toMatchObject({ + input: [{ role: "user", content: "hi" }], + metadata: { model: "llama-collected", provider: "groq" }, + }); + expect( + rows.find((row) => row.error?.includes("collected failure")), + ).toBeDefined(); + expect(diagnostic).toHaveBeenCalledWith( + expect.stringContaining("attempting collect-only tracing"), + expect.any(Error), + ); + }); + + it.each(["completed", "failed"] as const)( + "leaves a collect-only %s task pending when its child cannot be emitted", + async (status) => { + await traced( + async () => { + const originalLog = backgroundLogger.log.bind(backgroundLogger); + let logCalls = 0; + const logSpy = vi + .spyOn(backgroundLogger, "log") + .mockImplementation((items) => { + logCalls += 1; + if (logCalls > 1) { + throw new Error( + "simulated collect-only child emission failure", + ); + } + originalLog(items); + }); + await collectGroqBatchTrace({ + batch: { + id: "batch_collect_only_emission_failure", + endpoint: "/v1/chat/completions", + input_file_id: "file_collect_only_emission_failure", + status, + created_at: Date.now() / 1000, + completed_at: Date.now() / 1000 + 10, + failed_at: Date.now() / 1000 + 10, + request_counts: { completed: 1, failed: 0, total: 1 }, + }, + inputFile: jsonl([chatInput[0]]), + outputFile: jsonl([ + { + custom_id: "ok", + response: { status_code: 200, body: { choices: [] } }, + }, + ]), + }); + logSpy.mockRestore(); + }, + { name: "collect-only-emission-parent" }, + ); + + const rows = (await backgroundLogger.drain()) as Array< + Record + >; + const task = rows.find( + (row) => row.span_attributes?.name === "groq.batch", + ); + expect(task?.error).toBeUndefined(); + expect(task?.metrics?.end).toBeUndefined(); + }, + ); + + it("snapshots mutable collect-only result records before advancing the iterable", async () => { + const sharedRecord = { + custom_id: "ok", + response: { + status_code: 200, + body: { + choices: [ + { message: { role: "assistant", content: "first result" } }, + ], + }, + }, + }; + const outputFile = (function* () { + yield sharedRecord; + sharedRecord.custom_id = "bad"; + sharedRecord.response.body.choices[0].message.content = "second result"; + yield sharedRecord; + })(); + + await traced( + async () => { + await collectGroqBatchTrace({ + batch: { + id: "batch_mutable_collect_only", + endpoint: "/v1/chat/completions", + input_file_id: "file_mutable_collect_only", + status: "completed", + completed_at: Date.now() / 1000 + 10, + request_counts: { completed: 2, failed: 0, total: 2 }, + }, + inputFile: jsonl(chatInput), + outputFile, + }); + }, + { name: "mutable-collect-parent" }, + ); + + const rows = (await backgroundLogger.drain()) as Array>; + expect( + rows.find((row) => row.input?.[0]?.content === "hi")?.output?.[0]?.message + ?.content, + ).toBe("first result"); + expect( + rows.find((row) => row.input?.[0]?.content === "fail")?.output?.[0] + ?.message?.content, + ).toBe("second result"); + }); + + it("rejects malformed completed results before emitting any updates", async () => { + const { batch, input, traceContext } = await startBatch(); + await backgroundLogger.drain(); + await collectGroqBatchTrace({ + batch: { + ...batch, + status: "completed", + completed_at: Date.now() / 1000 + 10, + request_counts: { completed: 1, failed: 1, total: 2 }, + }, + traceContext, + inputFile: input, + outputFile: jsonl([ + { custom_id: "ok", response: { body: { choices: [] } } }, + ]), + errorFile: jsonl([ + { custom_id: "bad", error: { message: "valid failure" } }, + ]), + }); + + expect(await backgroundLogger.drain()).toEqual([]); + }); + + it("keeps the task pending when result totals disagree with request counts", async () => { + const { batch, input, traceContext } = await startBatch(); + await backgroundLogger.drain(); + await collectGroqBatchTrace({ + batch: { + ...batch, + status: "completed", + completed_at: Date.now() / 1000 + 10, + request_counts: { completed: 2, failed: 0, total: 2 }, + }, + traceContext, + inputFile: input, + outputFile: jsonl([ + { + custom_id: "ok", + response: { status_code: 200, body: { choices: [] } }, + }, + ]), + errorFile: jsonl([ + { custom_id: "bad", error: { message: "counted failure" } }, + ]), + }); + + expect(await backgroundLogger.drain()).toEqual([]); + }); + + it("skips inconsistent collect-only batches without creating pending spans", async () => { + await traced( + async () => { + await collectGroqBatchTrace({ + batch: { + id: "batch_collect_only_incomplete", + endpoint: "/v1/chat/completions", + input_file_id: "file_collect_only_incomplete", + status: "completed", + completed_at: Date.now() / 1000 + 10, + request_counts: { completed: 1, failed: 0, total: 2 }, + }, + inputFile: jsonl(chatInput), + outputFile: jsonl([ + { custom_id: "ok", response: { body: { choices: [] } } }, + ]), + }); + }, + { name: "incomplete-collect-parent" }, + ); + + const rows = (await backgroundLogger.drain()) as Array>; + expect(rows.map((row) => row.span_attributes?.name)).toEqual([ + "incomplete-collect-parent", + ]); + }); + + it("ends unresolved spans for failed batches despite inconsistent counts", async () => { + const { batch, input, traceContext } = await startBatch(); + await backgroundLogger.drain(); + await collectGroqBatchTrace({ + batch: { + ...batch, + status: "failed", + failed_at: Date.now() / 1000 + 10, + request_counts: { completed: 0, failed: 0, total: 1 }, + }, + traceContext, + inputFile: input, + outputFile: jsonl([ + { + custom_id: "ok", + response: { status_code: 200, body: { choices: [] } }, + }, + ]), + }); + + const rows = (await backgroundLogger.drain()) as Array>; + expect(rows).toHaveLength(3); + expect(rows.every((row) => row.metrics?.end !== undefined)).toBe(true); + expect( + rows.find((row) => row.span_attributes?.name === "groq.batch")?.error, + ).toContain("Groq Batch failed"); + expect( + rows.find((row) => row.error?.includes("Groq Batch failed")), + ).toBeDefined(); + }); + + it("fails every pending span after submission rejection", async () => { + const { params, input } = await startBatch(); + await backgroundLogger.drain(); + const error = new Error("create rejected"); + + await failGroqBatchTrace({ params, input, error }); + + const rows = (await backgroundLogger.drain()) as Array>; + expect(rows).toHaveLength(3); + expect(rows.every((row) => row.metrics?.end !== undefined)).toBe(true); + expect(rows.every((row) => row.error?.includes("create rejected"))).toBe( + true, + ); + }); + + it("replays iterable factories when a submission fails", async () => { + let factoryCalls = 0; + const input = () => { + factoryCalls += 1; + return chatInput.values(); + }; + const { params } = await startBatch(input); + await backgroundLogger.drain(); + + await failGroqBatchTrace({ + params, + input, + error: new Error("factory submission failed"), + }); + + const rows = (await backgroundLogger.drain()) as Array>; + expect(factoryCalls).toBeGreaterThanOrEqual(2); + expect(rows).toHaveLength(3); + expect(rows.every((row) => row.metrics?.end !== undefined)).toBe(true); + }); + + it("validates a failure replay before writing any span updates", async () => { + let mismatch = false; + const input = () => (mismatch ? [chatInput[0]] : chatInput); + const { params } = await startBatch(input); + await backgroundLogger.drain(); + mismatch = true; + + await failGroqBatchTrace({ + params, + input, + error: new Error("mismatched submission failed"), + }); + + expect(await backgroundLogger.drain()).toEqual([]); + }); + + it("traces valid requests while skipping record-local input errors", async () => { + const input = [ + chatInput[0], + { + custom_id: "invalid", + method: "GET", + url: "/v1/chat/completions", + body: {}, + }, + ]; + const params = await startGroqBatchTrace({ + inputFile: { id: "file_invalid" }, + input: () => input, + params: { + endpoint: "/v1/chat/completions", + completion_window: "24h", + }, + }); + + expect(params.metadata?.[CONTEXT_KEY]).toEqual(expect.any(String)); + const startedRows = (await backgroundLogger.drain()) as Array< + Record + >; + expect(startedRows).toHaveLength(2); + expect( + startedRows.find( + (row) => row.span_attributes?.name === "groq.chat.completions.create", + )?.input, + ).toEqual([{ role: "user", content: "hi" }]); + + const batch = { id: "batch_invalid_line", ...params, status: "validating" }; + const traceContext = await bindGroqBatchTrace({ batch }); + await collectGroqBatchTrace({ + batch: { + ...batch, + status: "completed", + completed_at: Date.now() / 1000 + 10, + request_counts: { completed: 1, failed: 0, total: 1 }, + }, + traceContext, + inputFile: () => input, + outputFile: jsonl([ + { + custom_id: "ok", + response: { status_code: 200, body: { choices: [] } }, + }, + ]), + }); + + const completedRows = (await backgroundLogger.drain()) as Array< + Record + >; + expect(completedRows).toHaveLength(2); + expect(completedRows.every((row) => row.metrics?.end !== undefined)).toBe( + true, + ); + }); + + it("supports a single input and result record beyond 16 MiB", async () => { + const largeContent = "x".repeat(17 * 1024 * 1024); + const input = [ + { + custom_id: "large-record", + method: "POST", + url: "/v1/chat/completions", + body: { + model: "llama-test", + messages: [{ role: "user", content: largeContent }], + }, + }, + ]; + const params = await startGroqBatchTrace({ + inputFile: { id: "file_large_record" }, + input: () => input, + params: { + endpoint: "/v1/chat/completions", + completion_window: "24h", + }, + }); + const batch = { id: "batch_large_record", ...params, status: "validating" }; + const traceContext = await bindGroqBatchTrace({ batch }); + + expect(params.metadata?.[CONTEXT_KEY]).toEqual(expect.any(String)); + expect(await backgroundLogger.drain()).toHaveLength(2); + + await collectGroqBatchTrace({ + batch: { + ...batch, + status: "completed", + completed_at: Date.now() / 1000 + 10, + request_counts: { completed: 1, failed: 0, total: 1 }, + }, + traceContext, + inputFile: () => input, + outputFile: () => [ + { + custom_id: "large-record", + response: { + status_code: 200, + body: { + choices: [ + { message: { role: "assistant", content: largeContent } }, + ], + }, + }, + }, + ], + }); + + const rows = (await backgroundLogger.drain()) as Array>; + expect(rows).toHaveLength(2); + expect(rows.every((row) => row.metrics?.end !== undefined)).toBe(true); + }); + + it("streams start inputs beyond the former 16 MiB staging limit", async () => { + const recordCount = 17; + let factoryCalls = 0; + let recordsRead = 0; + const content = "🙂".repeat(256 * 1024); + const largeInput = Array.from({ length: recordCount }, (_, index) => ({ + custom_id: `large-input-${index}`, + method: "POST", + url: "/v1/chat/completions", + body: { + model: "llama-test", + messages: [{ role: "user", content }], + }, + })); + const params = await startGroqBatchTrace({ + inputFile: { id: "file_non_ascii" }, + input: () => { + factoryCalls += 1; + return (function* () { + for (const record of largeInput) { + recordsRead += 1; + yield record; + } + })(); + }, + params: { + endpoint: "/v1/chat/completions", + completion_window: "24h", + }, + }); + + expect(params.metadata?.[CONTEXT_KEY]).toEqual(expect.any(String)); + expect(factoryCalls).toBeGreaterThanOrEqual(2); + expect(recordsRead).toBe(factoryCalls * recordCount); + const rows = (await backgroundLogger.drain()) as Array>; + expect(rows).toHaveLength(recordCount + 1); + expect( + rows.filter( + (row) => row.span_attributes?.name === "groq.chat.completions.create", + ), + ).toHaveLength(recordCount); + }); + + it("cancels response bodies when a staging limit stops consumption", async () => { + const cancel = vi.fn(); + const oversizedRecord = jsonl([ + { + ...chatInput[0], + body: { + ...chatInput[0].body, + messages: [ + { + role: "user", + content: "🙂".repeat(4 * 1024 * 1024), + }, + ], + }, + }, + ]); + const inputFile = new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode(`${oversizedRecord}\n`)); + }, + cancel, + }), + ); + + await traced( + async () => { + await collectGroqBatchTrace({ + batch: { + id: "batch_oversized_response", + endpoint: "/v1/chat/completions", + input_file_id: "file_oversized_response", + status: "failed", + failed_at: Date.now() / 1000, + }, + inputFile, + }); + }, + { name: "oversized-response-parent" }, + ); + + expect(cancel).toHaveBeenCalledOnce(); + expect(inputFile.bodyUsed).toBe(true); + await backgroundLogger.drain(); + }); + + it("preserves conflicting metadata without creating spans", async () => { + const params = await startGroqBatchTrace({ + inputFile: { id: "file_conflict" }, + input: () => [chatInput[0]], + params: { + endpoint: "/v1/chat/completions", + completion_window: "24h", + metadata: { [CONTEXT_KEY]: "caller-owned" }, + }, + }); + + expect(params.metadata).toEqual({ [CONTEXT_KEY]: "caller-owned" }); + expect(await backgroundLogger.drain()).toEqual([]); + }); + + it("does not include invalid file-source contents in diagnostics", async () => { + const diagnostic = vi.spyOn(debugLogger, "debug"); + const secretSource = { secret: "do-not-log-this-value" }; + + await collectGroqBatchTrace({ + batch: { + id: "batch_invalid_source", + endpoint: "/v1/chat/completions", + input_file_id: "file_invalid_source", + status: "completed", + }, + inputFile: secretSource as never, + outputFile: jsonl([]), + }); + + expect(diagnostic).toHaveBeenCalledWith( + expect.stringContaining("skipped invalid JSONL source"), + expect.any(Error), + ); + expect( + diagnostic.mock.calls.some((call) => call.includes(secretSource)), + ).toBe(false); + }); +}); diff --git a/js/src/instrumentation/plugins/groq-batch-instrumentation.ts b/js/src/instrumentation/plugins/groq-batch-instrumentation.ts new file mode 100644 index 000000000..122284bd0 --- /dev/null +++ b/js/src/instrumentation/plugins/groq-batch-instrumentation.ts @@ -0,0 +1,2786 @@ +import { debugLogger } from "../../debug-logger"; +import { + _internalGetGlobalState, + _internalStartSpanWithInitialMerge, + _internalStartSpanWithInitialMergeAndParentSpanIds, + flush, + getSpanParentObject, + NOOP_SPAN, + type Span, + utf8ByteLength, + withCurrent, +} from "../../logger"; +import { + INSTRUMENTATION_NAMES, + withSpanInstrumentationName, +} from "../../span-origin"; +import { getCurrentUnixTimestamp } from "../../util"; +import { isObject, SpanTypeAttribute } from "../../../util/index"; +import { + SpanComponentsV4, + type SpanComponentsV4Data, +} from "../../../util/span_identifier_v4"; +import type { + BindGroqBatchTraceArgs, + CollectGroqBatchTraceArgs, + FailGroqBatchTraceArgs, + GroqBatchCreateParams, + GroqBatchJSONL, + GroqBatchLike, + GroqBatchReplayableJSONL, + GroqBatchTraceContext, + StartGroqBatchTraceArgs, +} from "../../groq-batch-types"; +import { groqChannels } from "./groq-channels"; +import { extractGroqCompletionInput, parseGroqMetrics } from "./groq-span-data"; + +const BRAINTRUST_GROQ_BATCH_CONTEXT_KEY = "braintrust.batch_context"; +const BRAINTRUST_GROQ_BATCH_BINDING_KEY = "braintrust.batch_binding"; +const TERMINAL_STATUSES = new Set([ + "completed", + "failed", + "expired", + "cancelled", +]); +const INPUT_DIGEST_CHUNK_SIZE = 1024 * 1024; +const MAX_CONTEXT_LENGTH = 512; +const MAX_BATCH_INPUT_RECORDS = 50_000; +const MAX_BATCH_FILE_BYTES = 100 * 1024 * 1024; +const MAX_STAGED_INPUT_BYTES = 16 * 1024 * 1024; +const MAX_STAGED_RESULT_BYTES = 16 * 1024 * 1024; +const MAX_BATCH_RECORD_BYTES = MAX_BATCH_FILE_BYTES; +const MAX_BATCH_SPANS_BETWEEN_FLUSHES = 1000; + +type BatchContext = { + version: 1; + parent: string; + inputFileId: string; + inputDigest: string; + endpoint: "/v1/chat/completions"; + startTime: number; + operationNonce: string; + repairInputs: boolean; +}; + +type SerializedBatchContext = Omit & { + signature: string; +}; + +type BatchBinding = { + version: 1; + batchId: string; + operationNonce: string; +}; + +type SerializedBatchBinding = BatchBinding & { + signature: string; +}; + +type JSONValue = + | null + | boolean + | number + | string + | JSONValue[] + | { [key: string]: JSONValue }; + +type ParsedBatchRecord = { + value: JSONValue; + serialized: string; + encodedBytes: number; +}; + +type BatchInputRecord = { + customId: string; + spanData?: ReturnType; +}; + +type BatchResultRecord = { + customId: string; + error?: Error; + outcome: "completed" | "failed"; + responseBody?: Record; +}; + +type BatchFile = + | GroqBatchJSONL + | CollectGroqBatchTraceArgs["outputFile"]; + +type PreparedBatch = { + context: BatchContext; + input: GroqBatchReplayableJSONL; + inputData: BatchInputData; + params: GroqBatchCreateParams; +}; + +type BatchInputData = { + inputCount: number; + inputDigest: string; + inputs?: Map; + issues: Error[]; + remaining: Set; +}; + +type BatchResultCounts = { + completed: number; + failed: number; +}; + +type StagedBatchResult = { + input: BatchInputRecord; + result: BatchResultRecord; +}; + +type BatchFlushState = { + spansSinceFlush: number; +}; + +type BatchFileDigestState = { + chunk: string; + chunkBytes: number; + chunkDigests: string[]; +}; + +class BatchFatalError extends Error {} +class BatchRecordLimitError extends BatchFatalError {} +class BatchEmissionError extends Error {} + +function batchReplayIssue( + replay: BatchInputData, + expected: BatchInputData, + changedInputMessage: string, +): Error | undefined { + const inputIssue = replay.issues[0]; + if (inputIssue instanceof BatchFatalError) { + return inputIssue; + } + if ( + replay.inputCount !== expected.inputCount || + replay.inputDigest !== expected.inputDigest + ) { + return new Error(changedInputMessage); + } + return undefined; +} + +function read(value: unknown, key: PropertyKey): unknown { + if (!isObject(value)) { + return undefined; + } + try { + return Reflect.get(value, key); + } catch { + return undefined; + } +} + +function isBatchRecordIterable( + value: unknown, +): value is Iterable | AsyncIterable { + if ( + value === null || + (typeof value !== "object" && typeof value !== "function") + ) { + return false; + } + try { + return ( + typeof Reflect.get(value, Symbol.iterator) === "function" || + typeof Reflect.get(value, Symbol.asyncIterator) === "function" + ); + } catch { + return false; + } +} + +function replayableBatchInput(input: GroqBatchReplayableJSONL): GroqBatchJSONL { + if (typeof input === "string") { + return input; + } + if (typeof input === "function") { + const source = input(); + if (isBatchRecordIterable(source)) { + return source; + } + } + throw new Error( + "Groq Batch start and failure input must be a string or replayable iterable factory", + ); +} + +function validCustomId(value: unknown): value is string { + return typeof value === "string" && value.length > 0; +} + +function logBatchInstrumentationError(context: string, error: unknown): void { + debugLogger.debug(`Groq Batch instrumentation ${context}:`, error); +} + +function recordBatchIssue(issues: Error[], error: Error): void { + if ( + issues.length === 0 || + (error instanceof BatchFatalError && + !(issues[0] instanceof BatchFatalError)) + ) { + issues[0] = error; + } +} + +async function applyBatchWriteBackpressure( + state: BatchFlushState, +): Promise { + state.spansSinceFlush += 1; + if (state.spansSinceFlush >= MAX_BATCH_SPANS_BETWEEN_FLUSHES) { + await flush(); + state.spansSinceFlush = 0; + } +} + +function batchContextPayload(context: BatchContext): string { + return JSON.stringify(["groq", BRAINTRUST_GROQ_BATCH_CONTEXT_KEY, context]); +} + +function batchBindingPayload(binding: BatchBinding): string { + return JSON.stringify(["groq", BRAINTRUST_GROQ_BATCH_BINDING_KEY, binding]); +} + +async function importBatchSigningKey(secret: string) { + return await globalThis.crypto.subtle.importKey( + "raw", + new TextEncoder().encode(secret), + { name: "HMAC", hash: "SHA-256" }, + false, + ["sign", "verify"], + ); +} + +async function signBatchContext( + context: BatchContext, + secret: string, +): Promise { + try { + const signature = await globalThis.crypto.subtle.sign( + "HMAC", + await importBatchSigningKey(secret), + new TextEncoder().encode(batchContextPayload(context)), + ); + return digestHex(new Uint8Array(signature), 32); + } catch { + return undefined; + } +} + +async function signedBatchContextValue( + context: BatchContext, + secret: string, +): Promise { + const signature = await signBatchContext(context, secret); + if (!signature) { + return undefined; + } + const serialized = JSON.stringify({ + version: context.version, + parent: context.parent, + inputDigest: context.inputDigest, + startTime: context.startTime, + operationNonce: context.operationNonce, + repairInputs: context.repairInputs, + signature, + } satisfies SerializedBatchContext); + return serialized.length <= MAX_CONTEXT_LENGTH ? serialized : undefined; +} + +async function signBatchBinding( + binding: BatchBinding, + secret: string, +): Promise { + try { + const signature = await globalThis.crypto.subtle.sign( + "HMAC", + await importBatchSigningKey(secret), + new TextEncoder().encode(batchBindingPayload(binding)), + ); + return digestHex(new Uint8Array(signature), 32); + } catch { + return undefined; + } +} + +function signatureBytes(value: string): Uint8Array | undefined { + if (!/^[0-9a-f]{64}$/.test(value)) { + return undefined; + } + return Uint8Array.from(value.match(/.{2}/g) ?? [], (byte) => + Number.parseInt(byte, 16), + ); +} + +async function verifyBatchContext( + context: BatchContext, + signature: string, +): Promise { + const secret = + _internalGetGlobalState()._internalGetTraceContextSigningSecret(); + const bytes = signatureBytes(signature); + if (!secret || !bytes) { + return false; + } + try { + return await globalThis.crypto.subtle.verify( + "HMAC", + await importBatchSigningKey(secret), + bytes, + new TextEncoder().encode(batchContextPayload(context)), + ); + } catch { + return false; + } +} + +async function verifyBatchBinding( + binding: BatchBinding, + signature: string, +): Promise { + const secret = + _internalGetGlobalState()._internalGetTraceContextSigningSecret(); + const bytes = signatureBytes(signature); + if (!secret || !bytes) { + return false; + } + try { + return await globalThis.crypto.subtle.verify( + "HMAC", + await importBatchSigningKey(secret), + bytes, + new TextEncoder().encode(batchBindingPayload(binding)), + ); + } catch { + return false; + } +} + +async function parseBatchContext( + value: unknown, + inputFileId: string, + endpoint: string, +): Promise { + if ( + typeof value !== "string" || + endpoint !== "/v1/chat/completions" || + value.length > MAX_CONTEXT_LENGTH + ) { + return undefined; + } + + let parsed: unknown; + try { + parsed = JSON.parse(value); + } catch { + return undefined; + } + if ( + !isObject(parsed) || + parsed.version !== 1 || + typeof parsed.parent !== "string" || + typeof parsed.inputDigest !== "string" || + !/^[0-9a-f]{64}$/.test(parsed.inputDigest) || + typeof parsed.startTime !== "number" || + !Number.isFinite(parsed.startTime) || + typeof parsed.operationNonce !== "string" || + !/^[0-9a-f]{32}$/.test(parsed.operationNonce) || + typeof parsed.repairInputs !== "boolean" || + typeof parsed.signature !== "string" + ) { + return undefined; + } + + const context: BatchContext = { + version: 1, + parent: parsed.parent, + inputFileId, + inputDigest: parsed.inputDigest, + endpoint, + startTime: parsed.startTime, + operationNonce: parsed.operationNonce, + repairInputs: parsed.repairInputs, + }; + if (!(await verifyBatchContext(context, parsed.signature))) { + return undefined; + } + try { + SpanComponentsV4.fromStr(context.parent); + } catch { + return undefined; + } + return context; +} + +async function exportParent(parent: ReturnType) { + if ("toStr" in parent && typeof parent.toStr === "function") { + return parent.toStr(); + } + return await parent.export(); +} + +async function exportMinimalParent( + parent: ReturnType, +): Promise { + const exported = await exportParent(parent); + if (!exported) { + return undefined; + } + const parsed = SpanComponentsV4.fromStr(exported).data; + const projectId = read(parsed.compute_object_metadata_args, "project_id"); + const projectName = read(parsed.compute_object_metadata_args, "project_name"); + let computeObjectMetadataArgs: + | { project_id?: string; project_name?: string } + | undefined; + if (typeof projectId === "string" || typeof projectName === "string") { + computeObjectMetadataArgs = {}; + if (typeof projectId === "string") { + computeObjectMetadataArgs.project_id = projectId; + } + if (typeof projectName === "string") { + computeObjectMetadataArgs.project_name = projectName; + } + } + if (!parsed.object_id && !computeObjectMetadataArgs) { + return undefined; + } + + const routing = parsed.object_id + ? { object_id: parsed.object_id } + : { compute_object_metadata_args: computeObjectMetadataArgs ?? {} }; + // SpanComponents requires all three span identifiers together. The + // conditional below preserves that invariant, but TypeScript cannot infer + // it through the object spreads. + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions + return new SpanComponentsV4({ + object_type: parsed.object_type, + ...routing, + ...(parsed.row_id && parsed.span_id && parsed.root_span_id + ? { + row_id: parsed.row_id, + span_id: parsed.span_id, + root_span_id: parsed.root_span_id, + } + : {}), + } as SpanComponentsV4Data).toStr(); +} + +async function prepareCreateParams( + args: StartGroqBatchTraceArgs, + parent: ReturnType, + startTime: number, +): Promise { + const params: GroqBatchCreateParams = { + ...args.params, + input_file_id: args.inputFile.id, + }; + const endpoint = read(params, "endpoint"); + const inputFileId = read(params, "input_file_id"); + if ( + endpoint !== "/v1/chat/completions" || + typeof inputFileId !== "string" || + inputFileId.length === 0 + ) { + return undefined; + } + + const metadataValue = read(params, "metadata"); + if ( + metadataValue !== undefined && + metadataValue !== null && + !isObject(metadataValue) + ) { + logBatchInstrumentationError( + "skipped context injection", + "invalid metadata", + ); + return undefined; + } + const metadata = metadataValue ?? {}; + if ( + Object.prototype.hasOwnProperty.call( + metadata, + BRAINTRUST_GROQ_BATCH_CONTEXT_KEY, + ) + ) { + logBatchInstrumentationError( + "skipped context injection", + "conflicting reserved metadata", + ); + return undefined; + } + + const exportedParent = await exportMinimalParent(parent); + if (!exportedParent) { + logBatchInstrumentationError( + "skipped context injection", + "no routable Braintrust parent", + ); + return undefined; + } + const signingSecret = + _internalGetGlobalState()._internalGetTraceContextSigningSecret(); + if (!signingSecret) { + logBatchInstrumentationError( + "skipped context injection", + "no signing secret is available", + ); + return undefined; + } + + const inputData = await scanBatchInputs(replayableBatchInput(args.input)); + const inputIssue = inputData.issues[0]; + if (inputIssue instanceof BatchFatalError) { + logBatchInstrumentationError("skipped invalid input file", inputIssue); + return undefined; + } + if (inputIssue) { + logBatchInstrumentationError( + "skipped invalid input records while preparing valid requests", + inputIssue, + ); + } + const context: BatchContext = { + version: 1, + parent: exportedParent, + inputFileId, + inputDigest: inputData.inputDigest, + endpoint, + startTime, + operationNonce: randomBatchNonce(), + repairInputs: false, + }; + const serialized = await signedBatchContextValue(context, signingSecret); + if (!serialized) { + logBatchInstrumentationError( + "skipped context injection", + "could not serialize signed batch context", + ); + return undefined; + } + + return { + context, + input: args.input, + inputData, + params: { + ...params, + metadata: { + ...metadata, + [BRAINTRUST_GROQ_BATCH_CONTEXT_KEY]: serialized, + }, + }, + }; +} + +async function deterministicDigest(namespace: string, ...parts: string[]) { + const encoded = new TextEncoder().encode( + [namespace, ...parts].map((part) => `${part.length}:${part}`).join("\0"), + ); + return new Uint8Array( + await globalThis.crypto.subtle.digest("SHA-256", encoded), + ); +} + +async function addBatchFileDigestRecord( + state: BatchFileDigestState, + source: "output" | "error", + record: ParsedBatchRecord, +): Promise { + const { serialized, encodedBytes: serializedBytes } = record; + const framed = `${source.length}:${source}${serializedBytes}:${serialized}`; + const framedBytes = + String(source.length).length + + 1 + + source.length + + String(serializedBytes).length + + 1 + + serializedBytes; + if ( + state.chunkBytes > 0 && + state.chunkBytes + framedBytes > INPUT_DIGEST_CHUNK_SIZE + ) { + state.chunkDigests.push( + digestHex( + await deterministicDigest("groq:batch:result:chunk", state.chunk), + 32, + ), + ); + state.chunk = ""; + state.chunkBytes = 0; + } + state.chunk += framed; + state.chunkBytes += framedBytes; +} + +async function finishBatchFileDigest( + state: BatchFileDigestState, +): Promise { + state.chunkDigests.push( + digestHex( + await deterministicDigest("groq:batch:result:chunk", state.chunk), + 32, + ), + ); + return digestHex( + await deterministicDigest("groq:batch:result", state.chunkDigests.join("")), + 32, + ); +} + +function digestHex(bytes: Uint8Array, length: number): string { + return Array.from(bytes.slice(0, length)) + .map((byte) => byte.toString(16).padStart(2, "0")) + .join(""); +} + +function digestUuid(bytes: Uint8Array): string { + const hex = digestHex(bytes, 16); + return `${hex.slice(0, 8)}-${hex.slice(8, 12)}-${hex.slice(12, 16)}-${hex.slice(16, 20)}-${hex.slice(20)}`; +} + +function randomBatchNonce(): string { + return digestHex(globalThis.crypto.getRandomValues(new Uint8Array(16)), 16); +} + +async function batchSpanIds(operationNonce: string): Promise<{ + rowId: string; + spanId: string; + rootSpanId: string; +}> { + const [row, span, root] = await Promise.all([ + deterministicDigest("groq:batch:row", operationNonce), + deterministicDigest("groq:batch:span", operationNonce), + deterministicDigest("groq:batch:root", operationNonce), + ]); + return { + rowId: digestUuid(row), + spanId: digestHex(span, 8), + rootSpanId: digestHex(root, 16), + }; +} + +async function childSpanIds(operationNonce: string, customId: string) { + const [row, span] = await Promise.all([ + deterministicDigest("groq:batch:child:row", operationNonce, customId), + deterministicDigest("groq:batch:child:span", operationNonce, customId), + ]); + return { rowId: digestUuid(row), spanId: digestHex(span, 8) }; +} + +async function unboundBatchContextFromValues( + value: GroqBatchLike | GroqBatchCreateParams, +): Promise { + const metadata = read(value, "metadata"); + const inputFileId = read(value, "input_file_id"); + const endpoint = read(value, "endpoint"); + if ( + !isObject(metadata) || + typeof inputFileId !== "string" || + typeof endpoint !== "string" + ) { + return undefined; + } + return await parseBatchContext( + read(metadata, BRAINTRUST_GROQ_BATCH_CONTEXT_KEY), + inputFileId, + endpoint, + ); +} + +async function bindBatchTrace( + args: BindGroqBatchTraceArgs, +): Promise { + const context = await unboundBatchContextFromValues(args.batch); + const batchId = read(args.batch, "id"); + const signingSecret = + _internalGetGlobalState()._internalGetTraceContextSigningSecret(); + if (!context || !validCustomId(batchId) || !signingSecret) { + logBatchInstrumentationError( + "could not bind provider batch ID", + new Error( + "Groq Batch start context, batch ID, or signing key is invalid", + ), + ); + return undefined; + } + + const binding: BatchBinding = { + version: 1, + batchId, + operationNonce: context.operationNonce, + }; + const signature = await signBatchBinding(binding, signingSecret); + if (!signature) { + logBatchInstrumentationError( + "could not bind provider batch ID", + new Error("Groq Batch binding could not be signed"), + ); + return undefined; + } + const value = JSON.stringify({ + ...binding, + signature, + } satisfies SerializedBatchBinding); + if (value.length > MAX_CONTEXT_LENGTH) { + logBatchInstrumentationError( + "could not bind provider batch ID", + new Error("Groq Batch binding exceeds the instrumentation limit"), + ); + return undefined; + } + return { value }; +} + +async function validateBatchTraceBinding( + batch: GroqBatchLike, + context: BatchContext, + traceContext: GroqBatchTraceContext | undefined, +): Promise { + const value = read(traceContext, "value"); + const batchId = read(batch, "id"); + if ( + typeof value !== "string" || + value.length > MAX_CONTEXT_LENGTH || + !validCustomId(batchId) + ) { + return false; + } + let parsed: unknown; + try { + parsed = JSON.parse(value); + } catch { + return false; + } + if ( + !isObject(parsed) || + parsed.version !== 1 || + parsed.batchId !== batchId || + parsed.operationNonce !== context.operationNonce || + typeof parsed.signature !== "string" + ) { + return false; + } + return await verifyBatchBinding( + { + version: 1, + batchId, + operationNonce: context.operationNonce, + }, + parsed.signature, + ); +} + +export const interceptGroqBatchTraceBind: Parameters< + typeof groqChannels.batchesBindTrace.intercept +>[0] = async (target, thisArg, args) => { + await Reflect.apply(target, thisArg, args); + try { + const traceContext = await bindBatchTrace(args[0]); + return traceContext ? { traceContext } : {}; + } catch (error) { + logBatchInstrumentationError("could not bind provider batch ID", error); + return {}; + } +}; + +async function startBatchSpan(context: BatchContext): Promise { + const ids = await batchSpanIds(context.operationNonce); + const parent = SpanComponentsV4.fromStr(context.parent); + const parentSpanIds = + parent.data.span_id && parent.data.root_span_id + ? { + parentSpanIds: [parent.data.span_id], + rootSpanId: parent.data.root_span_id, + } + : { + parentSpanIds: [], + rootSpanId: ids.rootSpanId, + }; + return withCurrent(NOOP_SPAN, () => + _internalStartSpanWithInitialMergeAndParentSpanIds( + withSpanInstrumentationName( + { + name: "groq.batch", + type: SpanTypeAttribute.TASK, + parent: context.parent, + parentSpanIds, + spanId: ids.spanId, + startTime: context.startTime, + event: { + id: ids.rowId, + metadata: { provider: "groq" }, + }, + }, + INSTRUMENTATION_NAMES.GROQ, + ), + ), + ); +} + +async function startBatchChild( + context: BatchContext, + taskParent: string, + input: BatchInputRecord, +): Promise { + const ids = await childSpanIds(context.operationNonce, input.customId); + return withCurrent(NOOP_SPAN, () => + _internalStartSpanWithInitialMerge( + withSpanInstrumentationName( + { + name: "groq.chat.completions.create", + type: SpanTypeAttribute.LLM, + parent: taskParent, + spanId: ids.spanId, + startTime: context.startTime, + event: { + id: ids.rowId, + ...(input.spanData?.input !== undefined + ? { input: input.spanData.input } + : {}), + metadata: { + ...input.spanData?.metadata, + provider: "groq", + }, + }, + }, + INSTRUMENTATION_NAMES.GROQ, + ), + ), + ); +} + +async function* jsonlRecords( + file: BatchFile, + onIssue: (error: Error) => void = () => {}, +): AsyncGenerator { + let resolvedFile = await file; + if (typeof resolvedFile === "function") { + resolvedFile = await resolvedFile(); + } + if (resolvedFile === undefined) { + return; + } + if (typeof resolvedFile === "string") { + if (utf8ByteLength(resolvedFile) > MAX_BATCH_FILE_BYTES) { + onIssue(new BatchFatalError("Groq Batch file exceeds the 100 MB limit")); + return; + } + let offset = 0; + while (offset < resolvedFile.length) { + const newline = resolvedFile.indexOf("\n", offset); + const end = newline === -1 ? resolvedFile.length : newline; + const line = resolvedFile.slice(offset, end).replace(/\r$/, ""); + offset = newline === -1 ? resolvedFile.length : newline + 1; + if (!line.trim()) { + continue; + } + const lineBytes = utf8ByteLength(line); + if (lineBytes > MAX_BATCH_RECORD_BYTES) { + onIssue(new BatchRecordLimitError("Groq Batch record exceeds 100 MB")); + return; + } + try { + yield canonicalBatchRecord(JSON.parse(line)); + } catch (error) { + logBatchInstrumentationError("skipped malformed JSONL", error); + onIssue( + error instanceof BatchFatalError + ? error + : new Error("Groq Batch file contains malformed JSONL"), + ); + } + } + return; + } + + const body = read(resolvedFile, "body"); + const getReader = read(body, "getReader"); + if (isObject(body) && typeof getReader === "function") { + let reader: unknown; + let fullyConsumed = false; + try { + reader = Reflect.apply(getReader, body, []); + if (!isObject(reader)) { + throw new Error("Response body returned an invalid stream reader"); + } + const decoder = new TextDecoder(); + let encodedBytes = 0; + let pending = ""; + let pendingLineBytes = 0; + while (true) { + const readChunk = read(reader, "read"); + if (typeof readChunk !== "function") { + throw new Error("Response body stream reader has no read method"); + } + const chunk = await Reflect.apply(readChunk, reader, []); + if (!isObject(chunk)) { + throw new Error("Response body stream returned an invalid chunk"); + } + if (chunk.done === true) { + pending += decoder.decode(); + fullyConsumed = true; + break; + } + if (!(chunk.value instanceof Uint8Array)) { + throw new Error("Response body stream returned a non-byte chunk"); + } + encodedBytes += chunk.value.byteLength; + if (encodedBytes > MAX_BATCH_FILE_BYTES) { + onIssue( + new BatchFatalError("Groq Batch response exceeds the 100 MB limit"), + ); + return; + } + for (const byte of chunk.value) { + if (byte === 0x0a) { + pendingLineBytes = 0; + } else { + pendingLineBytes += 1; + if (pendingLineBytes > MAX_BATCH_RECORD_BYTES) { + onIssue( + new BatchRecordLimitError( + "Groq Batch response record exceeds 100 MB", + ), + ); + return; + } + } + } + const decoded = decoder.decode(chunk.value, { stream: true }); + pending += decoded; + let newline = pending.indexOf("\n"); + while (newline !== -1) { + const line = pending.slice(0, newline).replace(/\r$/, ""); + pending = pending.slice(newline + 1); + if (line.trim()) { + try { + yield canonicalBatchRecord(JSON.parse(line)); + } catch (error) { + logBatchInstrumentationError("skipped malformed JSONL", error); + onIssue( + error instanceof BatchFatalError + ? error + : new Error("Groq Batch response contains malformed JSONL"), + ); + } + } + newline = pending.indexOf("\n"); + } + } + if (pending.trim()) { + try { + yield canonicalBatchRecord(JSON.parse(pending.replace(/\r$/, ""))); + } catch (error) { + logBatchInstrumentationError("skipped malformed JSONL", error); + onIssue( + error instanceof BatchFatalError + ? error + : new Error("Groq Batch response contains malformed JSONL"), + ); + } + } + } catch (error) { + logBatchInstrumentationError("could not read JSONL response body", error); + onIssue( + new BatchFatalError("Groq Batch response body could not be read"), + ); + } finally { + const cancel = read(reader, "cancel"); + if (!fullyConsumed && isObject(reader) && typeof cancel === "function") { + try { + await Reflect.apply(cancel, reader, []); + } catch (error) { + logBatchInstrumentationError( + "could not cancel partially consumed response body", + error, + ); + } + } + const releaseLock = read(reader, "releaseLock"); + if (isObject(reader) && typeof releaseLock === "function") { + try { + Reflect.apply(releaseLock, reader, []); + } catch (error) { + logBatchInstrumentationError( + "could not release stream reader", + error, + ); + } + } + } + return; + } + + if (isBatchRecordIterable(resolvedFile)) { + let encodedBytes = 0; + let recordCount = 0; + for await (const record of resolvedFile) { + let normalizedRecord: ParsedBatchRecord; + try { + normalizedRecord = normalizeIterableBatchRecord(record); + } catch (error) { + logBatchInstrumentationError( + "could not measure iterable JSONL record", + error, + ); + onIssue( + error instanceof BatchRecordLimitError + ? error + : new Error("Groq Batch iterable contains a non-JSON record"), + ); + if (error instanceof BatchRecordLimitError) { + return; + } + continue; + } + encodedBytes += normalizedRecord.encodedBytes + (recordCount > 0 ? 1 : 0); + recordCount += 1; + if (encodedBytes > MAX_BATCH_FILE_BYTES) { + onIssue( + new BatchFatalError("Groq Batch iterable exceeds the 100 MB limit"), + ); + return; + } + yield normalizedRecord; + } + } else { + logBatchInstrumentationError( + "skipped invalid JSONL source", + new Error("Groq Batch file source has an unsupported type"), + ); + onIssue(new BatchFatalError("Groq Batch file source is invalid")); + } +} + +async function closePartiallyStartedBatch( + context: BatchContext, + task: Span, + taskParent: string, + attemptedCustomIds: Set, + error: Error, +): Promise { + const endTime = Math.max(getCurrentUnixTimestamp(), context.startTime); + const flushState: BatchFlushState = { spansSinceFlush: 0 }; + let childEmissionFailed = false; + for (const customId of attemptedCustomIds) { + try { + const child = await startBatchChild(context, taskParent, { customId }); + child.log({ error }); + child.end({ endTime }); + await applyBatchWriteBackpressure(flushState); + } catch (closeError) { + childEmissionFailed = true; + logBatchInstrumentationError( + "could not close partially started batch child", + closeError, + ); + } + } + if (childEmissionFailed) { + await flush(); + return false; + } + try { + task.log({ error }); + task.end({ endTime }); + await flush(); + return true; + } catch (closeError) { + logBatchInstrumentationError( + "could not close partially started batch task", + closeError, + ); + await flush(); + return false; + } +} + +async function startPreparedBatch(prepared: PreparedBatch): Promise { + const { context, input, inputData } = prepared; + const validationData = await scanBatchInputs(replayableBatchInput(input)); + const validationInputIssue = validationData.issues[0]; + if ( + validationInputIssue && + !(validationInputIssue instanceof BatchFatalError) + ) { + logBatchInstrumentationError( + "skipped invalid replay records while validating valid requests", + validationInputIssue, + ); + } + const validationIssue = batchReplayIssue( + validationData, + inputData, + "Groq Batch replayable input changed between passes", + ); + if (validationIssue) { + logBatchInstrumentationError( + "withheld correlation context before starting batch spans", + validationIssue, + ); + return false; + } + + await flush(); + const task = await startBatchSpan(context); + const taskParent = await task.export(); + await flush(); + const flushState: BatchFlushState = { spansSinceFlush: 0 }; + const attemptedCustomIds = new Set(); + let emissionFailed = false; + const replayData = await scanBatchInputs(replayableBatchInput(input), { + onInput: async (record) => { + if (emissionFailed) { + return; + } + attemptedCustomIds.add(record.customId); + try { + await startBatchChild(context, taskParent, record); + await applyBatchWriteBackpressure(flushState); + } catch (error) { + emissionFailed = true; + logBatchInstrumentationError("could not start batch child", error); + } + }, + }); + const replayInputIssue = replayData.issues[0]; + if (replayInputIssue && !(replayInputIssue instanceof BatchFatalError)) { + logBatchInstrumentationError( + "skipped invalid input records while starting valid requests", + replayInputIssue, + ); + } + const replayIssue = batchReplayIssue( + replayData, + inputData, + "Groq Batch replayable input changed between passes", + ); + if (replayIssue || emissionFailed) { + const startError = + replayIssue ?? + new BatchEmissionError("Groq Batch child span emission was incomplete"); + logBatchInstrumentationError( + "could not fully start batch spans", + startError, + ); + if ( + await closePartiallyStartedBatch( + context, + task, + taskParent, + attemptedCustomIds, + startError, + ) + ) { + return false; + } + logBatchInstrumentationError( + "preserved correlation context for partially started batch repair", + startError, + ); + const signingSecret = + _internalGetGlobalState()._internalGetTraceContextSigningSecret(); + const repairContext = { ...context, repairInputs: true }; + const repairValue = signingSecret + ? await signedBatchContextValue(repairContext, signingSecret) + : undefined; + if (repairValue) { + prepared.context = repairContext; + prepared.params = { + ...prepared.params, + metadata: { + ...prepared.params.metadata, + [BRAINTRUST_GROQ_BATCH_CONTEXT_KEY]: repairValue, + }, + }; + } else { + logBatchInstrumentationError( + "could not mark partial batch context for input repair", + new Error("Groq Batch repair context could not be signed"), + ); + } + return true; + } + await flush(); + return true; +} + +export const interceptGroqBatchTraceStart: Parameters< + typeof groqChannels.batchesStartTrace.intercept +>[0] = async (target, thisArg, args) => { + const parent = getSpanParentObject(); + const startTime = getCurrentUnixTimestamp(); + let prepared: PreparedBatch | undefined; + try { + prepared = await prepareCreateParams(args[0], parent, startTime); + } catch (error) { + logBatchInstrumentationError("could not prepare batch", error); + } + + let started = false; + if (prepared) { + try { + started = await startPreparedBatch(prepared); + } catch (error) { + logBatchInstrumentationError("could not start batch spans", error); + } + } + return await Reflect.apply(target, thisArg, [ + prepared && started ? { ...args[0], params: prepared.params } : args[0], + ]); +}; + +function serializedJSONStringBytes(value: string, maxBytes?: number): number { + let bytes = 2; + if (maxBytes !== undefined && bytes > maxBytes) { + throw new BatchRecordLimitError( + "Groq Batch record exceeds the 100 MB limit", + ); + } + for (let index = 0; index < value.length; index += 1) { + const code = value.charCodeAt(index); + if ( + code === 0x22 || + code === 0x5c || + code === 0x08 || + code === 0x09 || + code === 0x0a || + code === 0x0c || + code === 0x0d + ) { + bytes += 2; + } else if (code <= 0x1f) { + bytes += 6; + } else if (code >= 0xd800 && code <= 0xdbff) { + const next = value.charCodeAt(index + 1); + if (next >= 0xdc00 && next <= 0xdfff) { + bytes += 4; + index += 1; + } else { + bytes += 6; + } + } else if (code >= 0xdc00 && code <= 0xdfff) { + bytes += 6; + } else if (code <= 0x7f) { + bytes += 1; + } else if (code <= 0x7ff) { + bytes += 2; + } else { + bytes += 3; + } + if (maxBytes !== undefined && bytes > maxBytes) { + throw new BatchRecordLimitError( + "Groq Batch record exceeds the 100 MB limit", + ); + } + } + return bytes; +} + +function canonicalizeJSON(value: unknown, maxBytes?: number): JSONValue { + const byteBudget = { used: 0, max: maxBytes }; + const ancestors = new Set(); + + function consumeBytes(serialized: string): void { + byteBudget.used += utf8ByteLength(serialized); + if (byteBudget.max !== undefined && byteBudget.used > byteBudget.max) { + throw new BatchRecordLimitError( + "Groq Batch record exceeds the 100 MB limit", + ); + } + } + + function consumeJSONStringBytes(stringValue: string): void { + const remaining = + byteBudget.max === undefined + ? undefined + : byteBudget.max - byteBudget.used; + byteBudget.used += serializedJSONStringBytes(stringValue, remaining); + } + + function visit(entry: unknown): JSONValue { + if (entry === null) { + consumeBytes("null"); + return entry; + } + if (typeof entry === "string" || typeof entry === "boolean") { + if (typeof entry === "string") { + consumeJSONStringBytes(entry); + } else { + consumeBytes(JSON.stringify(entry)); + } + return entry; + } + if (typeof entry === "number") { + if (!Number.isFinite(entry)) { + throw new Error("Groq Batch input contains a non-finite number"); + } + consumeBytes(JSON.stringify(entry)); + return entry; + } + if (Array.isArray(entry)) { + if (ancestors.has(entry)) { + throw new Error("Groq Batch input contains a circular value"); + } + ancestors.add(entry); + consumeBytes("[]"); + try { + const result: JSONValue[] = []; + for (let index = 0; index < entry.length; index += 1) { + if (index > 0) { + consumeBytes(","); + } + result.push(visit(entry[index])); + } + return result; + } finally { + ancestors.delete(entry); + } + } + if (!isObject(entry)) { + throw new Error("Groq Batch input contains a non-JSON value"); + } + if (ancestors.has(entry)) { + throw new Error("Groq Batch input contains a circular value"); + } + ancestors.add(entry); + consumeBytes("{}"); + try { + const result: Array<[string, JSONValue]> = []; + const keys = Object.keys(entry).sort(); + for (let index = 0; index < keys.length; index += 1) { + const key = keys[index]; + const item = read(entry, key); + if (item === undefined) { + throw new Error("Groq Batch input contains an undefined value"); + } + if (index > 0) { + consumeBytes(","); + } + consumeJSONStringBytes(key); + consumeBytes(":"); + result.push([key, visit(item)]); + } + return Object.fromEntries(result); + } finally { + ancestors.delete(entry); + } + } + + return visit(value); +} + +function canonicalBatchRecord( + value: unknown, + maxBytes = MAX_BATCH_RECORD_BYTES, +): ParsedBatchRecord { + const canonicalValue = canonicalizeJSON(value, maxBytes); + const serialized = JSON.stringify(canonicalValue); + if (serialized === undefined) { + throw new Error("Groq Batch record could not be serialized"); + } + return { + value: canonicalValue, + serialized, + encodedBytes: utf8ByteLength(serialized), + }; +} + +function normalizeIterableBatchRecord(value: unknown): ParsedBatchRecord { + let estimatedBytes = 0; + const serialized = JSON.stringify(value, function (key, entry) { + const arrayEntry = Array.isArray(this); + const omitted = + !arrayEntry && + key !== "" && + (entry === undefined || + typeof entry === "function" || + typeof entry === "symbol"); + if (omitted) { + return entry; + } + + if (key !== "") { + estimatedBytes += 1; + if (!arrayEntry) { + estimatedBytes += + serializedJSONStringBytes( + key, + MAX_BATCH_RECORD_BYTES - estimatedBytes, + ) + 1; + } + } + if ( + entry === undefined || + typeof entry === "function" || + typeof entry === "symbol" || + (typeof entry === "number" && !Number.isFinite(entry)) + ) { + estimatedBytes += 4; + } else if (typeof entry === "string") { + estimatedBytes += serializedJSONStringBytes( + entry, + MAX_BATCH_RECORD_BYTES - estimatedBytes, + ); + } else if (entry === null) { + estimatedBytes += 4; + } else if (typeof entry === "boolean") { + estimatedBytes += entry ? 4 : 5; + } else if (typeof entry === "number") { + estimatedBytes += String(entry).length; + } else if (typeof entry === "object") { + estimatedBytes += 2; + } + if (estimatedBytes > MAX_BATCH_RECORD_BYTES) { + throw new BatchRecordLimitError( + "Groq Batch record exceeds the 100 MB limit", + ); + } + return entry; + }); + if (typeof serialized !== "string") { + throw new Error("Groq Batch iterable contains a non-JSON record"); + } + const standardBytes = utf8ByteLength(serialized); + if (standardBytes > MAX_BATCH_RECORD_BYTES) { + throw new BatchRecordLimitError( + "Groq Batch record exceeds the 100 MB limit", + ); + } + return canonicalBatchRecord(JSON.parse(serialized)); +} + +async function scanBatchInputs( + file: BatchFile, + options: { + retainInputs?: boolean; + allowUnstagedInputs?: boolean; + onInput?: (input: BatchInputRecord) => Promise; + } = {}, +): Promise { + let inputs = options.retainInputs + ? new Map() + : undefined; + const issues: Error[] = []; + const chunkDigests: string[] = []; + const remaining = new Set(); + let bufferedInputBytes = 0; + let digestChunk = ""; + let digestChunkBytes = 0; + let totalCanonicalBytes = 0; + try { + for await (const record of jsonlRecords(file, (issue) => + recordBatchIssue(issues, issue), + )) { + const value = record.value; + const customId = read(value, "custom_id"); + const body = read(value, "body"); + if ( + !validCustomId(customId) || + remaining.has(customId) || + read(value, "method") !== "POST" || + read(value, "url") !== "/v1/chat/completions" || + !isObject(body) + ) { + recordBatchIssue( + issues, + new Error("Groq Batch input contains an invalid record"), + ); + continue; + } + if (remaining.size >= MAX_BATCH_INPUT_RECORDS) { + recordBatchIssue( + issues, + new BatchFatalError( + `Groq Batch input exceeds the ${MAX_BATCH_INPUT_RECORDS} record limit`, + ), + ); + break; + } + + let canonicalRecord: ParsedBatchRecord; + try { + canonicalRecord = canonicalBatchRecord({ + body: canonicalizeJSON(body), + custom_id: customId, + }); + } catch (error) { + logBatchInstrumentationError("skipped invalid input record", error); + recordBatchIssue( + issues, + new Error("Groq Batch input contains a non-JSON record"), + ); + continue; + } + const canonicalRecordBytes = canonicalRecord.encodedBytes; + if ( + inputs && + bufferedInputBytes + canonicalRecordBytes > MAX_STAGED_INPUT_BYTES + ) { + inputs.clear(); + inputs = undefined; + if (!options.allowUnstagedInputs) { + recordBatchIssue( + issues, + new BatchFatalError( + "Groq Batch staged input exceeds the 16 MiB buffer", + ), + ); + break; + } + } + totalCanonicalBytes += canonicalRecordBytes; + if (totalCanonicalBytes > MAX_BATCH_FILE_BYTES) { + recordBatchIssue( + issues, + new BatchFatalError("Groq Batch input exceeds the 100 MB limit"), + ); + break; + } + const framedRecord = `${canonicalRecordBytes}:${canonicalRecord.serialized}`; + const framedRecordBytes = + String(canonicalRecordBytes).length + 1 + canonicalRecordBytes; + if ( + digestChunkBytes > 0 && + digestChunkBytes + framedRecordBytes > INPUT_DIGEST_CHUNK_SIZE + ) { + chunkDigests.push( + digestHex( + await deterministicDigest("groq:batch:input:chunk", digestChunk), + 32, + ), + ); + digestChunk = ""; + digestChunkBytes = 0; + } + digestChunk += framedRecord; + digestChunkBytes += framedRecordBytes; + remaining.add(customId); + + let spanData: BatchInputRecord["spanData"]; + if (inputs || options.onInput) { + try { + spanData = extractGroqCompletionInput(body); + } catch (error) { + logBatchInstrumentationError("could not extract batch input", error); + } + } + const input = { customId, spanData }; + if (inputs) { + bufferedInputBytes += canonicalRecordBytes; + inputs.set(customId, input); + } + await options.onInput?.(input); + } + } catch (error) { + logBatchInstrumentationError("could not process input file", error); + recordBatchIssue( + issues, + new BatchFatalError("Groq Batch input file could not be processed"), + ); + } + chunkDigests.push( + digestHex( + await deterministicDigest("groq:batch:input:chunk", digestChunk), + 32, + ), + ); + return { + inputCount: remaining.size, + inputDigest: digestHex( + await deterministicDigest("groq:batch:input", chunkDigests.join("")), + 32, + ), + inputs, + issues, + remaining, + }; +} + +function batchResultError(value: unknown): Error | undefined { + if (typeof value === "string" && value.length > 0) { + return new Error(value); + } + if (!isObject(value)) { + return undefined; + } + const message = read(value, "message"); + return new Error( + typeof message === "string" && message.length > 0 + ? message + : "Groq Batch request failed", + ); +} + +function validateBatchResult( + value: unknown, + source: "output" | "error", +): BatchResultRecord | undefined { + if (!isObject(value)) { + return undefined; + } + const customId = read(value, "custom_id"); + if (!validCustomId(customId)) { + return undefined; + } + + const response = read(value, "response"); + const statusCode = read(response, "status_code"); + const responseBody = read(response, "body"); + const topLevelError = batchResultError(read(value, "error")); + if (source === "output") { + if ( + !isObject(response) || + typeof statusCode !== "number" || + !Number.isInteger(statusCode) || + statusCode < 200 || + statusCode >= 300 || + !isObject(responseBody) || + !Array.isArray(read(responseBody, "choices")) || + topLevelError + ) { + return undefined; + } + return { + customId, + outcome: "completed", + responseBody, + }; + } + + const responseError = batchResultError(read(responseBody, "error")); + const hasFailedStatus = + typeof statusCode === "number" && + Number.isInteger(statusCode) && + (statusCode < 200 || statusCode >= 300); + if (!topLevelError && !responseError && !hasFailedStatus) { + return undefined; + } + return { + customId, + error: + topLevelError ?? responseError ?? new Error("Groq Batch request failed"), + outcome: "failed", + }; +} + +async function completeBatchResult( + context: BatchContext, + taskParent: string, + endTime: number, + input: BatchInputRecord, + result: BatchResultRecord, +): Promise { + const child = await startBatchChild(context, taskParent, input); + if (result.error) { + child.log({ error: result.error }); + } else if (result.responseBody) { + const model = read(result.responseBody, "model"); + child.log({ + output: read(result.responseBody, "choices"), + ...(typeof model === "string" ? { metadata: { model } } : {}), + metrics: parseGroqMetrics(result.responseBody), + }); + } + child.end({ endTime }); +} + +async function completeResultFile({ + context, + endTime, + file, + inputData, + issues, + resultCounts, + seen, + source, + taskParent, + flushState, + digestState, +}: { + context: BatchContext; + endTime: number; + file: BatchFile; + inputData: BatchInputData; + issues: Error[]; + resultCounts: BatchResultCounts; + seen: Set; + source: "output" | "error"; + taskParent: string; + flushState: BatchFlushState; + digestState?: BatchFileDigestState; +}): Promise { + let resultRecordCount = 0; + let emissionSucceeded = true; + try { + for await (const record of jsonlRecords(file, (issue) => + recordBatchIssue(issues, issue), + )) { + const value = record.value; + try { + resultRecordCount += 1; + if (resultRecordCount > MAX_BATCH_INPUT_RECORDS) { + recordBatchIssue( + issues, + new Error( + `Groq Batch result exceeds the ${MAX_BATCH_INPUT_RECORDS} record limit`, + ), + ); + break; + } + if (digestState) { + await addBatchFileDigestRecord(digestState, source, record); + } + const result = validateBatchResult(value, source); + if (!result) { + recordBatchIssue( + issues, + new Error("Groq Batch result contains an invalid record"), + ); + continue; + } + if (seen.has(result.customId)) { + recordBatchIssue( + issues, + new Error("Groq Batch result contains a duplicate"), + ); + continue; + } + if (!inputData.remaining.has(result.customId)) { + recordBatchIssue( + issues, + new Error("Groq Batch result does not match an input record"), + ); + continue; + } + const input = inputData.inputs?.get(result.customId) ?? { + customId: result.customId, + }; + await completeBatchResult(context, taskParent, endTime, input, result); + await applyBatchWriteBackpressure(flushState); + seen.add(result.customId); + inputData.remaining.delete(result.customId); + resultCounts[result.outcome] += 1; + } catch (error) { + emissionSucceeded = false; + logBatchInstrumentationError("skipped invalid result", error); + recordBatchIssue( + issues, + new BatchEmissionError( + "Groq Batch child span emission was incomplete", + ), + ); + } + } + } catch (error) { + logBatchInstrumentationError(`could not process ${source} file`, error); + recordBatchIssue( + issues, + new Error(`Groq Batch ${source} file could not be processed`), + ); + } + return emissionSucceeded; +} + +async function stageResultFile({ + file, + inputData, + issues, + resultCounts, + seen, + source, + stagedBytes, + stagedResults, + digestState, +}: { + file: BatchFile; + inputData: BatchInputData; + issues: Error[]; + resultCounts: BatchResultCounts; + seen: Set; + source: "output" | "error"; + stagedBytes?: { value: number }; + stagedResults?: Map; + digestState?: BatchFileDigestState; +}): Promise { + let resultRecordCount = 0; + try { + for await (const record of jsonlRecords(file, (issue) => + recordBatchIssue(issues, issue), + )) { + const value = record.value; + resultRecordCount += 1; + if (resultRecordCount > MAX_BATCH_INPUT_RECORDS) { + recordBatchIssue( + issues, + new Error( + `Groq Batch result exceeds the ${MAX_BATCH_INPUT_RECORDS} record limit`, + ), + ); + return false; + } + + if (digestState) { + await addBatchFileDigestRecord(digestState, source, record); + } + + const result = validateBatchResult(value, source); + if (!result) { + recordBatchIssue( + issues, + new Error("Groq Batch result contains an invalid record"), + ); + continue; + } + if (seen.has(result.customId)) { + recordBatchIssue( + issues, + new Error("Groq Batch result contains a duplicate"), + ); + continue; + } + if (!inputData.remaining.has(result.customId)) { + recordBatchIssue( + issues, + new Error("Groq Batch result does not match an input record"), + ); + continue; + } + const input = inputData.inputs?.get(result.customId) ?? { + customId: result.customId, + }; + if (stagedBytes) { + stagedBytes.value += record.encodedBytes; + } + if (stagedBytes && stagedBytes.value > MAX_STAGED_RESULT_BYTES) { + recordBatchIssue( + issues, + new Error("Groq Batch staged results exceed the 16 MiB buffer"), + ); + return false; + } + + seen.add(result.customId); + inputData.remaining.delete(result.customId); + resultCounts[result.outcome] += 1; + stagedResults?.set(result.customId, { input, result }); + } + } catch (error) { + logBatchInstrumentationError(`could not stage ${source} file`, error); + recordBatchIssue( + issues, + new Error(`Groq Batch ${source} file could not be staged`), + ); + } + return true; +} + +function terminalEndTime(batch: GroqBatchLike, startTime: number): number { + let timestamp: number | null | undefined; + switch (batch.status) { + case "completed": + timestamp = batch.completed_at; + break; + case "failed": + timestamp = batch.failed_at; + break; + case "expired": + timestamp = batch.expired_at; + break; + default: + timestamp = batch.cancelled_at; + } + return typeof timestamp === "number" && + Number.isFinite(timestamp) && + timestamp >= startTime + ? timestamp + : Math.max(getCurrentUnixTimestamp(), startTime); +} + +async function validatedBatchInputs( + file: BatchFile, + context: BatchContext | undefined, + retainInputs = false, +): Promise { + const inputData = await scanBatchInputs(file, { + retainInputs, + allowUnstagedInputs: context?.repairInputs, + }); + const issues = [...inputData.issues]; + const inputIssue = issues[0]; + if ( + context && + !(inputIssue instanceof BatchFatalError) && + inputData.inputDigest !== context.inputDigest + ) { + recordBatchIssue( + issues, + new BatchFatalError( + "Groq Batch input does not match the signed input file", + ), + ); + } + return { ...inputData, issues }; +} + +async function collectOnlyBatchContext( + batch: GroqBatchLike, + inputData: BatchInputData, + parent: ReturnType, + collectionTime: number, +): Promise { + const exportedParent = await exportParent(parent); + const batchId = read(batch, "id"); + const inputFileId = read(batch, "input_file_id"); + const endpoint = read(batch, "endpoint"); + if ( + !exportedParent || + typeof batchId !== "string" || + batchId.length === 0 || + typeof inputFileId !== "string" || + inputFileId.length === 0 || + endpoint !== "/v1/chat/completions" + ) { + return undefined; + } + const createdAt = read(batch, "created_at"); + const startTime = + typeof createdAt === "number" && + Number.isFinite(createdAt) && + createdAt >= 0 && + createdAt <= collectionTime + ? createdAt + : collectionTime; + return { + version: 1, + parent: exportedParent, + inputFileId, + inputDigest: inputData.inputDigest, + endpoint, + startTime, + repairInputs: false, + operationNonce: digestHex( + await deterministicDigest( + "groq:batch:collect-only", + exportedParent, + batchId, + inputData.inputDigest, + ), + 16, + ), + }; +} + +function validateBatchResultCounts( + batch: GroqBatchLike, + inputCount: number, + actual: BatchResultCounts, +): Error[] { + const requestCounts = read(batch, "request_counts"); + if (requestCounts === undefined || requestCounts === null) { + return []; + } + if (!isObject(requestCounts)) { + return [new Error("Groq Batch request counts are invalid")]; + } + + const issues: Error[] = []; + for (const [key, expected, observed] of [ + ["total", read(requestCounts, "total"), inputCount], + ["completed", read(requestCounts, "completed"), actual.completed], + ["failed", read(requestCounts, "failed"), actual.failed], + ] as const) { + if ( + expected !== undefined && + (typeof expected !== "number" || + !Number.isInteger(expected) || + expected < 0 || + expected !== observed) + ) { + issues.push(new Error(`Groq Batch ${key} request count is inconsistent`)); + } + } + return issues; +} + +function completedBatchResultSourcesIssue( + args: CollectGroqBatchTraceArgs, +): Error | undefined { + const requestCounts = read(args.batch, "request_counts"); + if (requestCounts === undefined || requestCounts === null) { + return undefined; + } + if (!isObject(requestCounts)) { + return new Error("Groq Batch request counts are invalid"); + } + const completed = read(requestCounts, "completed"); + const failed = read(requestCounts, "failed"); + const total = read(requestCounts, "total"); + for (const [key, value] of [ + ["completed", completed], + ["failed", failed], + ["total", total], + ] as const) { + if ( + value !== undefined && + (typeof value !== "number" || !Number.isInteger(value) || value < 0) + ) { + return new Error(`Groq Batch ${key} request count is invalid`); + } + } + + const completedCount = typeof completed === "number" ? completed : undefined; + const failedCount = typeof failed === "number" ? failed : undefined; + const totalCount = typeof total === "number" ? total : undefined; + if (totalCount === 0) { + return undefined; + } + const requiresOutput = + completedCount !== undefined + ? completedCount > 0 + : totalCount !== undefined && + failedCount !== undefined && + failedCount < totalCount; + const requiresError = + failedCount !== undefined + ? failedCount > 0 + : totalCount !== undefined && + completedCount !== undefined && + completedCount < totalCount; + if (requiresOutput && args.outputFile === undefined) { + return new Error( + "Groq Batch completed collection requires the output file", + ); + } + if (requiresError && args.errorFile === undefined) { + return new Error("Groq Batch completed collection requires the error file"); + } + return undefined; +} + +async function resultFilesAreReplayable( + args: CollectGroqBatchTraceArgs, +): Promise { + for (const file of [args.outputFile, args.errorFile]) { + if (file === undefined) { + continue; + } + const resolved = await file; + if (typeof resolved !== "string" && typeof resolved !== "function") { + return false; + } + } + return true; +} + +async function completeReplayableSignedCompletedBatch( + args: CollectGroqBatchTraceArgs, + context: BatchContext, + inputData: BatchInputData, +): Promise { + const validationInputData = { + ...inputData, + remaining: new Set(inputData.remaining), + }; + const validationIssues: Error[] = []; + const validationCounts: BatchResultCounts = { completed: 0, failed: 0 }; + const validationDigest: BatchFileDigestState = { + chunk: "", + chunkBytes: 0, + chunkDigests: [], + }; + const validationSeen = new Set(); + const outputValid = await stageResultFile({ + file: args.outputFile, + inputData: validationInputData, + issues: validationIssues, + resultCounts: validationCounts, + seen: validationSeen, + source: "output", + digestState: validationDigest, + }); + const errorValid = await stageResultFile({ + file: args.errorFile, + inputData: validationInputData, + issues: validationIssues, + resultCounts: validationCounts, + seen: validationSeen, + source: "error", + digestState: validationDigest, + }); + const validationCountIssues = validateBatchResultCounts( + args.batch, + inputData.inputCount, + validationCounts, + ); + const validationIssue = validationIssues[0] ?? validationCountIssues[0]; + if (validationIssue instanceof BatchRecordLimitError) { + logBatchInstrumentationError( + "left completed batch spans pending after result record overflow", + validationIssue, + ); + return; + } + if ( + !outputValid || + !errorValid || + validationIssue || + validationInputData.remaining.size > 0 + ) { + logBatchInstrumentationError( + "left completed batch spans pending after replayable validation", + validationIssue ?? new Error("Groq Batch result files are incomplete"), + ); + return; + } + const expectedDigest = await finishBatchFileDigest(validationDigest); + + const endTime = terminalEndTime(args.batch, context.startTime); + await flush(); + const task = await startBatchSpan(context); + const taskParent = await task.export(); + await flush(); + const flushState: BatchFlushState = { spansSinceFlush: 0 }; + const emissionInputData = { + ...inputData, + remaining: new Set(inputData.remaining), + }; + const emissionIssues: Error[] = []; + const emissionCounts: BatchResultCounts = { completed: 0, failed: 0 }; + const emissionDigest: BatchFileDigestState = { + chunk: "", + chunkBytes: 0, + chunkDigests: [], + }; + const emissionSeen = new Set(); + const outputEmissionSucceeded = await completeResultFile({ + context, + endTime, + file: args.outputFile, + inputData: emissionInputData, + issues: emissionIssues, + resultCounts: emissionCounts, + seen: emissionSeen, + source: "output", + taskParent, + flushState, + digestState: emissionDigest, + }); + const errorEmissionSucceeded = await completeResultFile({ + context, + endTime, + file: args.errorFile, + inputData: emissionInputData, + issues: emissionIssues, + resultCounts: emissionCounts, + seen: emissionSeen, + source: "error", + taskParent, + flushState, + digestState: emissionDigest, + }); + const observedDigest = await finishBatchFileDigest(emissionDigest); + const emissionCountIssues = validateBatchResultCounts( + args.batch, + inputData.inputCount, + emissionCounts, + ); + const emissionIssue = emissionIssues[0] ?? emissionCountIssues[0]; + let collectionError: Error | undefined; + if (!outputEmissionSucceeded || !errorEmissionSucceeded) { + collectionError = new BatchEmissionError( + "Groq Batch child span emission was incomplete", + ); + } else if (emissionIssue) { + collectionError = emissionIssue; + } else if (observedDigest !== expectedDigest) { + collectionError = new Error( + "Groq Batch replayable result files changed between passes", + ); + } else if (emissionInputData.remaining.size > 0) { + collectionError = new Error( + "Groq Batch replay left unresolved child spans", + ); + } + if (collectionError) { + logBatchInstrumentationError( + "left completed batch retryable after emission failure", + collectionError, + ); + await flush(); + return; + } + task.end({ endTime }); + await flush(); +} + +async function completeSignedCompletedBatch( + args: CollectGroqBatchTraceArgs, + context: BatchContext, + inputData: BatchInputData, +): Promise { + if (await resultFilesAreReplayable(args)) { + await completeReplayableSignedCompletedBatch(args, context, inputData); + return; + } + + const seen = new Set(); + const issues: Error[] = []; + const resultCounts: BatchResultCounts = { completed: 0, failed: 0 }; + const stagedBytes = { value: 0 }; + const stagedResults = new Map(); + const outputWithinBounds = await stageResultFile({ + file: args.outputFile, + inputData, + issues, + resultCounts, + seen, + source: "output", + stagedBytes, + stagedResults, + }); + const errorWithinBounds = await stageResultFile({ + file: args.errorFile, + inputData, + issues, + resultCounts, + seen, + source: "error", + stagedBytes, + stagedResults, + }); + const countIssues = validateBatchResultCounts( + args.batch, + inputData.inputCount, + resultCounts, + ); + const consistencyIssue = issues[0] ?? countIssues[0]; + if (!outputWithinBounds || !errorWithinBounds) { + logBatchInstrumentationError( + "left completed batch retryable after one-shot staging overflow", + consistencyIssue ?? + new Error( + "Groq Batch result files exceeded the 16 MiB one-shot collection buffer", + ), + ); + return; + } + if (consistencyIssue instanceof BatchRecordLimitError) { + logBatchInstrumentationError( + "left completed batch retryable after result record overflow", + consistencyIssue, + ); + return; + } + if (consistencyIssue || inputData.remaining.size > 0) { + logBatchInstrumentationError( + "left completed batch spans pending without emitting partial results", + consistencyIssue ?? + new Error("Groq Batch result files are incomplete or exceed limits"), + ); + return; + } + + const endTime = terminalEndTime(args.batch, context.startTime); + await flush(); + const task = await startBatchSpan(context); + const taskParent = await task.export(); + await flush(); + const flushState: BatchFlushState = { spansSinceFlush: 0 }; + let emissionFailed = false; + for (const staged of stagedResults.values()) { + try { + await completeBatchResult( + context, + taskParent, + endTime, + staged.input, + staged.result, + ); + await applyBatchWriteBackpressure(flushState); + } catch (error) { + emissionFailed = true; + logBatchInstrumentationError("could not emit staged result", error); + } + } + if (emissionFailed) { + logBatchInstrumentationError( + "left completed batch retryable after child emission failure", + new Error("Groq Batch child span emission was incomplete"), + ); + await flush(); + return; + } + task.end({ endTime }); + await flush(); +} + +async function completeSignedBatch( + args: CollectGroqBatchTraceArgs, + context: BatchContext, + inputData: BatchInputData, + status: string, +): Promise { + if (status === "completed") { + await completeSignedCompletedBatch(args, context, inputData); + return; + } + + const endTime = terminalEndTime(args.batch, context.startTime); + await flush(); + const task = await startBatchSpan(context); + const taskParent = await task.export(); + await flush(); + const flushState: BatchFlushState = { spansSinceFlush: 0 }; + const seen = new Set(); + const issues: Error[] = []; + const resultCounts: BatchResultCounts = { completed: 0, failed: 0 }; + const outputEmissionSucceeded = await completeResultFile({ + context, + endTime, + file: args.outputFile, + inputData, + issues, + resultCounts, + seen, + source: "output", + taskParent, + flushState, + }); + const errorEmissionSucceeded = await completeResultFile({ + context, + endTime, + file: args.errorFile, + inputData, + issues, + resultCounts, + seen, + source: "error", + taskParent, + flushState, + }); + await flush(); + if (!outputEmissionSucceeded || !errorEmissionSucceeded) { + logBatchInstrumentationError( + `left ${status} batch retryable after child emission failure`, + new BatchEmissionError("Groq Batch child span emission was incomplete"), + ); + return; + } + const countIssues = validateBatchResultCounts( + args.batch, + inputData.inputCount, + resultCounts, + ); + const consistencyIssue = issues[0] ?? countIssues[0]; + if (consistencyIssue) { + logBatchInstrumentationError( + `completed ${status} batch with inconsistent results`, + consistencyIssue, + ); + } + let unresolvedEmissionFailed = false; + for (const customId of inputData.remaining) { + try { + const input = inputData.inputs?.get(customId) ?? { customId }; + const child = await startBatchChild(context, taskParent, input); + child.log({ error: new Error(`Groq Batch ${status}`) }); + child.end({ endTime }); + await applyBatchWriteBackpressure(flushState); + } catch (error) { + unresolvedEmissionFailed = true; + logBatchInstrumentationError("could not complete missing result", error); + } + } + if (unresolvedEmissionFailed) { + logBatchInstrumentationError( + `left ${status} batch retryable after unresolved child emission failure`, + new BatchEmissionError("Groq Batch child span emission was incomplete"), + ); + await flush(); + return; + } + task.log({ error: new Error(`Groq Batch ${status}`) }); + task.end({ endTime }); + await flush(); +} + +async function completeCollectOnlyBatch( + args: CollectGroqBatchTraceArgs, + context: BatchContext, + inputData: BatchInputData, + status: string, +): Promise { + const seen = new Set(); + const issues: Error[] = []; + const resultCounts: BatchResultCounts = { completed: 0, failed: 0 }; + const stagedBytes = { value: 0 }; + const stagedResults = new Map(); + const outputWithinBounds = await stageResultFile({ + file: args.outputFile, + inputData, + issues, + resultCounts, + seen, + source: "output", + stagedBytes, + stagedResults, + }); + const errorWithinBounds = await stageResultFile({ + file: args.errorFile, + inputData, + issues, + resultCounts, + seen, + source: "error", + stagedBytes, + stagedResults, + }); + if (!outputWithinBounds || !errorWithinBounds) { + logBatchInstrumentationError( + "skipped collect-only tracing", + issues[0] ?? new Error("Groq Batch results exceeded staging limits"), + ); + return; + } + + const countIssues = validateBatchResultCounts( + args.batch, + inputData.inputCount, + resultCounts, + ); + const consistencyIssue = issues[0] ?? countIssues[0]; + if ( + status === "completed" && + (consistencyIssue || inputData.remaining.size > 0) + ) { + logBatchInstrumentationError( + "skipped inconsistent collect-only tracing", + consistencyIssue ?? new Error("Groq Batch result files are incomplete"), + ); + return; + } + if (consistencyIssue) { + logBatchInstrumentationError( + `completed collect-only ${status} batch with inconsistent results`, + consistencyIssue, + ); + } + + const endTime = terminalEndTime(args.batch, context.startTime); + await flush(); + const task = await startBatchSpan(context); + const taskParent = await task.export(); + await flush(); + const flushState: BatchFlushState = { spansSinceFlush: 0 }; + let emissionFailed = false; + for (const input of inputData.inputs?.values() ?? []) { + try { + const staged = stagedResults.get(input.customId); + if (staged) { + await completeBatchResult( + context, + taskParent, + endTime, + staged.input, + staged.result, + ); + } else { + const child = await startBatchChild(context, taskParent, input); + child.log({ error: new Error(`Groq Batch ${status}`) }); + child.end({ endTime }); + } + await applyBatchWriteBackpressure(flushState); + } catch (error) { + emissionFailed = true; + logBatchInstrumentationError("could not emit collect-only child", error); + } + } + if (emissionFailed) { + logBatchInstrumentationError( + `left collect-only ${status} batch retryable after child emission failure`, + new BatchEmissionError("Groq Batch child span emission was incomplete"), + ); + await flush(); + return; + } + if (status !== "completed") { + task.log({ error: new Error(`Groq Batch ${status}`) }); + } + task.end({ endTime }); + await flush(); +} + +async function repairSignedBatchInputs( + file: BatchFile, + context: BatchContext, + expected: BatchInputData, +): Promise { + const validationData = await scanBatchInputs(file); + const validationIssue = validationData.issues[0]; + if ( + validationIssue instanceof BatchFatalError || + validationData.inputCount !== expected.inputCount || + validationData.inputDigest !== expected.inputDigest + ) { + logBatchInstrumentationError( + "left partial batch pending before input repair", + validationIssue ?? + new Error("Groq Batch replayable repair input changed between passes"), + ); + return false; + } + + await flush(); + const task = await startBatchSpan(context); + const taskParent = await task.export(); + await flush(); + const flushState: BatchFlushState = { spansSinceFlush: 0 }; + let emissionFailed = false; + const emissionData = await scanBatchInputs(file, { + onInput: async (input) => { + if (emissionFailed) { + return; + } + try { + await startBatchChild(context, taskParent, input); + await applyBatchWriteBackpressure(flushState); + } catch (error) { + emissionFailed = true; + logBatchInstrumentationError( + "could not repair partially started batch child input", + error, + ); + } + }, + }); + const emissionIssue = emissionData.issues[0]; + if ( + emissionFailed || + emissionIssue instanceof BatchFatalError || + emissionData.inputCount !== expected.inputCount || + emissionData.inputDigest !== expected.inputDigest + ) { + logBatchInstrumentationError( + "left partial batch pending after input repair failure", + emissionIssue ?? + new BatchEmissionError( + "Groq Batch child input repair emission was incomplete", + ), + ); + await flush(); + return false; + } + await flush(); + return true; +} + +async function collectBatch( + args: CollectGroqBatchTraceArgs, + collectParent: ReturnType, +): Promise { + const status = read(args.batch, "status"); + if (typeof status !== "string" || !TERMINAL_STATUSES.has(status)) { + return; + } + const startedContext = await unboundBatchContextFromValues(args.batch); + let signedContext: BatchContext | undefined; + if (startedContext) { + if ( + !(await validateBatchTraceBinding( + args.batch, + startedContext, + args.traceContext, + )) + ) { + logBatchInstrumentationError( + "left started batch spans pending", + new Error("Groq Batch provider ID binding is absent or invalid"), + ); + return; + } + signedContext = startedContext; + } else { + logBatchInstrumentationError( + "could not resume signed context; attempting collect-only tracing", + new Error("Groq Batch signed context is absent or invalid"), + ); + } + if (status === "completed") { + const sourceIssue = completedBatchResultSourcesIssue(args); + if (sourceIssue) { + logBatchInstrumentationError( + signedContext + ? "left completed batch spans pending without consuming partial files" + : "skipped collect-only tracing without consuming partial files", + sourceIssue, + ); + return; + } + } + const inputData = await validatedBatchInputs( + args.inputFile, + signedContext, + !signedContext || signedContext.repairInputs, + ); + const inputIssue = inputData.issues[0]; + if ( + inputIssue instanceof BatchFatalError || + (!signedContext && !inputData.inputs) + ) { + logBatchInstrumentationError( + signedContext + ? "left batch spans pending" + : "could not establish collect-only correlation", + inputIssue ?? + new Error("Groq Batch collect-only input could not be staged"), + ); + return; + } + if (inputIssue) { + logBatchInstrumentationError( + "skipped invalid input records while collecting valid requests", + inputIssue, + ); + } + + if (signedContext) { + if ( + signedContext.repairInputs && + !inputData.inputs && + !(await repairSignedBatchInputs(args.inputFile, signedContext, inputData)) + ) { + return; + } + await completeSignedBatch(args, signedContext, inputData, status); + return; + } + const context = await collectOnlyBatchContext( + args.batch, + inputData, + collectParent, + getCurrentUnixTimestamp(), + ); + if (!context) { + logBatchInstrumentationError( + "could not establish collect-only correlation", + new Error("Groq Batch has no routable active Braintrust parent"), + ); + return; + } + await completeCollectOnlyBatch(args, context, inputData, status); +} + +export const interceptGroqBatchTraceCollect: Parameters< + typeof groqChannels.batchesCollectTrace.intercept +>[0] = (target, thisArg, args) => { + const collectParent = getSpanParentObject(); + return Promise.resolve(Reflect.apply(target, thisArg, args)).then( + async (batch) => { + try { + await collectBatch({ ...args[0], batch }, collectParent); + } catch (error) { + logBatchInstrumentationError("could not collect batch", error); + } + return batch; + }, + ); +}; + +function submissionError(value: unknown): Error { + if (value instanceof Error) { + return value; + } + return new Error( + typeof value === "string" ? value : "Groq Batch submission failed", + ); +} + +async function failBatch(args: FailGroqBatchTraceArgs): Promise { + const context = await unboundBatchContextFromValues(args.params); + if (!context) { + logBatchInstrumentationError( + "could not resume submission failure", + new Error("Groq Batch signed context is absent or invalid"), + ); + return; + } + const inputData = await validatedBatchInputs( + replayableBatchInput(args.input), + context, + ); + const inputIssue = inputData.issues[0]; + if (inputIssue instanceof BatchFatalError) { + logBatchInstrumentationError("left batch spans pending", inputIssue); + return; + } + if (inputIssue) { + logBatchInstrumentationError( + "skipped invalid input records while failing valid requests", + inputIssue, + ); + } + const error = submissionError(args.error); + const endTime = Math.max(getCurrentUnixTimestamp(), context.startTime); + await flush(); + const task = await startBatchSpan(context); + const taskParent = await task.export(); + await flush(); + const flushState: BatchFlushState = { spansSinceFlush: 0 }; + let emissionFailed = false; + const replayData = await scanBatchInputs(replayableBatchInput(args.input), { + onInput: async (input) => { + try { + const child = await startBatchChild(context, taskParent, input); + child.log({ error }); + child.end({ endTime }); + await applyBatchWriteBackpressure(flushState); + } catch (instrumentationError) { + emissionFailed = true; + logBatchInstrumentationError( + "could not fail pending child", + instrumentationError, + ); + } + }, + }); + const replayInputIssue = replayData.issues[0]; + if (replayInputIssue && !(replayInputIssue instanceof BatchFatalError)) { + logBatchInstrumentationError( + "skipped invalid replay records while failing valid requests", + replayInputIssue, + ); + } + const replayIssue = batchReplayIssue( + replayData, + inputData, + "Groq Batch replayable failure input changed between passes", + ); + if (replayIssue || emissionFailed) { + logBatchInstrumentationError( + "left submission failure retryable", + replayIssue ?? new Error("Groq Batch child span emission was incomplete"), + ); + await flush(); + return; + } + task.log({ error }); + task.end({ endTime }); + await flush(); +} + +export const interceptGroqBatchTraceFail: Parameters< + typeof groqChannels.batchesFailTrace.intercept +>[0] = async (target, thisArg, args) => { + const result = await Reflect.apply(target, thisArg, args); + try { + await failBatch(args[0]); + } catch (error) { + logBatchInstrumentationError("could not fail batch", error); + } + return result; +}; diff --git a/js/src/instrumentation/plugins/groq-channels.ts b/js/src/instrumentation/plugins/groq-channels.ts index 7ed21de89..c4ab3414d 100644 --- a/js/src/instrumentation/plugins/groq-channels.ts +++ b/js/src/instrumentation/plugins/groq-channels.ts @@ -8,12 +8,50 @@ import type { GroqEmbeddingCreateParams, GroqEmbeddingResponse, } from "../../vendor-sdk-types/groq"; +import type { + BindGroqBatchTraceArgs, + CollectGroqBatchTraceArgs, + FailGroqBatchTraceArgs, + GroqBatchCreateParams, + GroqBatchLike, + GroqBatchTraceContext, + StartGroqBatchTraceArgs, +} from "../../groq-batch-types"; type GroqChatResult = GroqChatCompletion | GroqChatStream; export const groqChannels = defineChannels( "groq-sdk", { + batchesStartTrace: channel< + [StartGroqBatchTraceArgs], + GroqBatchCreateParams + >({ + channelName: "batches.start-trace", + kind: "async", + }), + + batchesBindTrace: channel< + [BindGroqBatchTraceArgs], + { traceContext?: GroqBatchTraceContext } + >({ + channelName: "batches.bind-trace", + kind: "async", + }), + + batchesCollectTrace: channel< + [CollectGroqBatchTraceArgs], + GroqBatchLike + >({ + channelName: "batches.collect-trace", + kind: "async", + }), + + batchesFailTrace: channel<[FailGroqBatchTraceArgs], void>({ + channelName: "batches.fail-trace", + kind: "async", + }), + chatCompletionsCreate: channel< [GroqChatCreateParams], GroqChatResult, diff --git a/js/src/instrumentation/plugins/groq-plugin.test.ts b/js/src/instrumentation/plugins/groq-plugin.test.ts index a140deb80..04e57e658 100644 --- a/js/src/instrumentation/plugins/groq-plugin.test.ts +++ b/js/src/instrumentation/plugins/groq-plugin.test.ts @@ -1,11 +1,33 @@ import { describe, expect, it } from "vitest"; -import { - aggregateGroqChatCompletionChunks, - parseGroqMetrics, -} from "./groq-plugin"; +import { aggregateGroqChatCompletionChunks } from "./groq-plugin"; +import { extractGroqCompletionInput, parseGroqMetrics } from "./groq-span-data"; + +describe("Groq completion input extraction", () => { + it("uses the same metadata allowlist for direct and batch calls", () => { + const params = { + messages: [{ role: "user", content: "hello" }], + model: "llama-test", + temperature: 0, + max_completion_tokens: 32, + user: "sensitive-user-id", + unexpected: "do-not-capture", + }; + const expected = { + input: [{ role: "user", content: "hello" }], + metadata: { + max_completion_tokens: 32, + model: "llama-test", + provider: "groq", + temperature: 0, + }, + }; + + expect(extractGroqCompletionInput(params)).toEqual(expected); + }); +}); describe("parseGroqMetrics", () => { - it("merges OpenAI-compatible usage metrics with Groq cache metrics", () => { + it("normalizes Groq cache usage into prompt_cached_tokens", () => { expect( parseGroqMetrics({ usage: { @@ -22,57 +44,156 @@ describe("parseGroqMetrics", () => { }), ).toEqual({ completion_tokens: 4, - dram_cached_tokens: 2, + prompt_cached_tokens: 5, prompt_tokens: 10, - sram_cached_tokens: 3, tokens: 14, }); }); + it("prefers canonical cached tokens over Groq hardware counters", () => { + expect( + parseGroqMetrics({ + usage: { + prompt_tokens: 10, + prompt_tokens_details: { cached_tokens: 4 }, + }, + x_groq: { + usage: { + dram_cached_tokens: 2, + sram_cached_tokens: 3, + }, + }, + }), + ).toEqual({ + prompt_cached_tokens: 4, + prompt_tokens: 10, + }); + }); + + it("omits Groq hardware cache usage that exceeds prompt tokens", () => { + expect( + parseGroqMetrics({ + usage: { prompt_tokens: 4 }, + x_groq: { + usage: { + dram_cached_tokens: 3, + sram_cached_tokens: 2, + }, + }, + }), + ).toEqual({ prompt_tokens: 4 }); + }); + + it("omits reasoning usage that exceeds completion tokens", () => { + expect( + parseGroqMetrics({ + usage: { + completion_tokens: 3, + completion_tokens_details: { reasoning_tokens: 4 }, + }, + }), + ).toEqual({ completion_tokens: 3 }); + }); + it("returns an empty object for unknown values", () => { expect(parseGroqMetrics(undefined)).toEqual({}); expect(parseGroqMetrics(null)).toEqual({}); expect(parseGroqMetrics({})).toEqual({}); }); + + it("preserves a reported zero-token cache hit", () => { + expect( + parseGroqMetrics({ + usage: { prompt_tokens: 0 }, + x_groq: { + usage: { dram_cached_tokens: 0, sram_cached_tokens: 0 }, + }, + }), + ).toEqual({ prompt_cached_tokens: 0, prompt_tokens: 0 }); + }); + + it("selects and validates only canonical token metrics", () => { + expect( + parseGroqMetrics({ + usage: { + prompt_tokens: -1, + input_tokens: 7, + completion_tokens: 2.5, + output_tokens: 3, + total_tokens: Number.POSITIVE_INFINITY, + queue_time: 0.1, + prompt_time: 0.2, + total_time: 0.3, + completion_tokens_details: { reasoning_tokens: 2 }, + }, + x_groq: { + usage: { + dram_cached_tokens: -4, + sram_cached_tokens: 1, + }, + }, + }), + ).toEqual({ + completion_reasoning_tokens: 2, + completion_tokens: 3, + prompt_cached_tokens: 1, + prompt_tokens: 7, + tokens: 10, + }); + }); }); describe("aggregateGroqChatCompletionChunks", () => { + it("preserves shared stream-result metrics", () => { + const aggregated = aggregateGroqChatCompletionChunks([], { + __braintrust_cached_metric: 1, + }); + + expect(aggregated.metrics).toEqual({ cached: 1 }); + }); + it("preserves parsed reasoning chunks", () => { - expect( - aggregateGroqChatCompletionChunks([ - { - choices: [ - { - delta: { - role: "assistant", - reasoning: "First, count the marbles. ", - }, - finish_reason: null, + const aggregated = aggregateGroqChatCompletionChunks([ + { + choices: [ + { + delta: { + role: "assistant", + reasoning: "First, count the marbles. ", }, - ], - }, - { - choices: [ - { - delta: { - reasoning: "Then double the remainder.", - }, - finish_reason: null, + finish_reason: null, + }, + ], + }, + { + choices: [ + { + delta: { + reasoning: "Then double the remainder.", }, - ], - }, - { - choices: [ - { - delta: { - content: "14", - }, - finish_reason: "stop", + finish_reason: null, + }, + ], + }, + { + choices: [ + { + delta: { + content: "14", }, - ], + finish_reason: "stop", + }, + ], + usage: { + prompt_tokens: 8, + completion_tokens: 4, + total_tokens: 12, + queue_time: 0.5, }, - ]).output, - ).toEqual([ + }, + ]); + + expect(aggregated.output).toEqual([ { finish_reason: "stop", index: 0, @@ -85,5 +206,10 @@ describe("aggregateGroqChatCompletionChunks", () => { }, }, ]); + expect(aggregated.metrics).toEqual({ + completion_tokens: 4, + prompt_tokens: 8, + tokens: 12, + }); }); }); diff --git a/js/src/instrumentation/plugins/groq-plugin.ts b/js/src/instrumentation/plugins/groq-plugin.ts index 10b6d2c0a..aa3941539 100644 --- a/js/src/instrumentation/plugins/groq-plugin.ts +++ b/js/src/instrumentation/plugins/groq-plugin.ts @@ -5,13 +5,16 @@ import { 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 { aggregateChatCompletionChunks } from "./openai-plugin"; import { groqChannels } from "./groq-channels"; +import { extractGroqCompletionInput, parseGroqMetrics } from "./groq-span-data"; +import { + interceptGroqBatchTraceBind, + interceptGroqBatchTraceCollect, + interceptGroqBatchTraceFail, + interceptGroqBatchTraceStart, +} from "./groq-batch-instrumentation"; import type { GroqChatCompletion, GroqChatCompletionChunk, @@ -19,17 +22,20 @@ import type { export class GroqPlugin extends BasePlugin { protected onEnable(): void { + this.unsubscribers.push( + groqChannels.batchesStartTrace.intercept(interceptGroqBatchTraceStart), + groqChannels.batchesBindTrace.intercept(interceptGroqBatchTraceBind), + groqChannels.batchesCollectTrace.intercept( + interceptGroqBatchTraceCollect, + ), + groqChannels.batchesFailTrace.intercept(interceptGroqBatchTraceFail), + ); + 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" }, - }; - }, + extractInput: ([params]) => extractGroqCompletionInput(params), extractOutput: (result) => result?.choices, extractMetrics: (result, startTime) => { const metrics = parseGroqMetrics(result); @@ -69,44 +75,6 @@ export class GroqPlugin extends BasePlugin { } } -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, @@ -120,6 +88,20 @@ export function aggregateGroqChatCompletionChunks( streamResult, endEvent, ); + const metrics: Record = { ...aggregated.metrics }; + const retainedMetricNames = new Set(["cached"]); + for (const chunk of chunks) { + const normalizedMetrics = parseGroqMetrics(chunk); + for (const name of Object.keys(normalizedMetrics)) { + retainedMetricNames.add(name); + } + Object.assign(metrics, normalizedMetrics); + } + for (const name of Object.keys(metrics)) { + if (!retainedMetricNames.has(name)) { + delete metrics[name]; + } + } const reasoning = aggregateGroqReasoning(chunks); if (reasoning !== undefined) { const message = aggregated.output[0]?.message; @@ -128,7 +110,7 @@ export function aggregateGroqChatCompletionChunks( } } return { - metrics: aggregated.metrics, + metrics, output: aggregated.output, }; } diff --git a/js/src/instrumentation/plugins/groq-span-data.ts b/js/src/instrumentation/plugins/groq-span-data.ts new file mode 100644 index 000000000..bc8cbebc4 --- /dev/null +++ b/js/src/instrumentation/plugins/groq-span-data.ts @@ -0,0 +1,175 @@ +import { processInputAttachments } from "../../wrappers/attachment-utils"; +import type { GroqChatCompletion } from "../../vendor-sdk-types/groq"; +import { isObject } from "../../../util/index"; + +const GROQ_COMPLETION_METADATA_KEYS = [ + "model", + "temperature", + "top_p", + "max_tokens", + "max_completion_tokens", + "frequency_penalty", + "presence_penalty", + "stop", + "response_format", + "tools", + "tool_choice", + "parallel_tool_calls", + "max_tool_calls", + "reasoning_format", + "reasoning_effort", + "seed", +] as const; + +export function extractGroqCompletionInput(params: Record): { + input: unknown; + metadata: Record; +} { + const metadata: Record = { provider: "groq" }; + for (const key of GROQ_COMPLETION_METADATA_KEYS) { + try { + const value = Reflect.get(params, key); + if (value !== undefined) { + metadata[key] = value; + } + } catch { + // Ignore hostile getters in provider input. + } + } + let messages: unknown; + try { + messages = Reflect.get(params, "messages"); + } catch { + // Ignore hostile getters in provider input. + } + return { + input: processInputAttachments(messages), + metadata, + }; +} + +function readProperty(value: unknown, key: PropertyKey): unknown { + if (!isObject(value)) { + return undefined; + } + try { + return Reflect.get(value, key); + } catch { + return undefined; + } +} + +function readNonnegativeInteger( + value: unknown, + key: PropertyKey, +): number | undefined { + const result = readProperty(value, key); + return typeof result === "number" && + Number.isSafeInteger(result) && + result >= 0 + ? result + : undefined; +} + +export function parseGroqMetrics( + result: + | Pick + | { usage?: unknown; x_groq?: unknown } + | null + | undefined, +): Record { + const metrics: Record = {}; + const usage = readProperty(result, "usage"); + const xGroq = readProperty(result, "x_groq"); + + for (const [source, target] of [ + ["prompt_tokens", "prompt_tokens"], + ["input_tokens", "prompt_tokens"], + ["completion_tokens", "completion_tokens"], + ["output_tokens", "completion_tokens"], + ["total_tokens", "tokens"], + ] as const) { + const value = readNonnegativeInteger(usage, source); + if (value !== undefined && metrics[target] === undefined) { + metrics[target] = value; + } + } + + if ( + metrics.tokens === undefined && + metrics.prompt_tokens !== undefined && + metrics.completion_tokens !== undefined + ) { + const total = metrics.prompt_tokens + metrics.completion_tokens; + if (Number.isSafeInteger(total)) { + metrics.tokens = total; + } + } + + const promptDetails = + readProperty(usage, "prompt_tokens_details") ?? + readProperty(usage, "input_tokens_details"); + const completionDetails = + readProperty(usage, "completion_tokens_details") ?? + readProperty(usage, "output_tokens_details"); + const promptCachedTokens = readNonnegativeInteger( + promptDetails, + "cached_tokens", + ); + const completionReasoningTokens = readNonnegativeInteger( + completionDetails, + "reasoning_tokens", + ); + if ( + promptCachedTokens !== undefined && + (metrics.prompt_tokens === undefined || + promptCachedTokens <= metrics.prompt_tokens) + ) { + metrics.prompt_cached_tokens = promptCachedTokens; + } + if ( + completionReasoningTokens !== undefined && + (metrics.completion_tokens === undefined || + completionReasoningTokens <= metrics.completion_tokens) + ) { + metrics.completion_reasoning_tokens = completionReasoningTokens; + } + + if (!isObject(xGroq)) { + return metrics; + } + + const extraUsage = readProperty(xGroq, "usage"); + if (!isObject(extraUsage)) { + return metrics; + } + + const dramCachedTokens = readNonnegativeInteger( + extraUsage, + "dram_cached_tokens", + ); + const sramCachedTokens = readNonnegativeInteger( + extraUsage, + "sram_cached_tokens", + ); + let groqCachedTokens: number | undefined; + for (const value of [dramCachedTokens, sramCachedTokens]) { + if (value !== undefined) { + const total = (groqCachedTokens ?? 0) + value; + if (!Number.isSafeInteger(total)) { + return metrics; + } + groqCachedTokens = total; + } + } + + if ( + metrics.prompt_cached_tokens === undefined && + groqCachedTokens !== undefined && + metrics.prompt_tokens !== undefined && + groqCachedTokens <= metrics.prompt_tokens + ) { + metrics.prompt_cached_tokens = groqCachedTokens; + } + return metrics; +} diff --git a/js/src/logger.ts b/js/src/logger.ts index ad316c481..1085e58cb 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -6732,6 +6732,22 @@ export function _internalStartSpanWithInitialMerge< InitialSpanWriteAsMergeArg).span; } +/** @internal Start a deterministic span under an exported object-only parent. */ +export function _internalStartSpanWithInitialMergeAndParentSpanIds< + IsAsyncFlush extends boolean = true, +>(args: StartSpanArgs & AsyncFlushArg & OptionalStateArg): Span { + return startSpanAndIsLogger( + { + ...args, + [INITIAL_SPAN_WRITE_AS_MERGE]: true, + } as StartSpanArgs & + AsyncFlushArg & + OptionalStateArg & + InitialSpanWriteAsMergeArg, + { useParentSpanIdsForObjectParent: true }, + ).span; +} + /** @internal Start a span with SDK-controlled context fields. */ export function _internalStartSpanWithContext< IsAsyncFlush extends boolean = true, @@ -6768,6 +6784,7 @@ function startSpanAndIsLogger( AsyncFlushArg & OptionalStateArg & InternalSpanContextArg, + internalOptions?: { useParentSpanIdsForObjectParent?: boolean }, ): { span: Span; isSyncFlushLogger: boolean } { const state = args?.state ?? _globalState; @@ -6785,17 +6802,21 @@ function startSpanAndIsLogger( parentObject instanceof SpanComponentsV3 || parentObject instanceof SpanComponentsV4 ) { - const parentSpanIds: ParentSpanIds | undefined = + let parentSpanIds: ParentSpanIds | MultiParentSpanIds | undefined; + if ( parentObject.data.row_id && parentSpanIdsUsable( parentObject.data.span_id, parentObject.data.root_span_id, ) - ? { - spanId: parentObject.data.span_id, - rootSpanId: parentObject.data.root_span_id, - } - : undefined; + ) { + parentSpanIds = { + spanId: parentObject.data.span_id, + rootSpanId: parentObject.data.root_span_id, + }; + } else if (internalOptions?.useParentSpanIdsForObjectParent) { + parentSpanIds = args?.parentSpanIds; + } // The parent object/state are already resolved from `parent` above; drop // the raw `parent` so it isn't re-normalized. const { parent: _ignoredParent, ...spanArgs } = args ?? {}; diff --git a/js/src/wrappers/groq.test.ts b/js/src/wrappers/groq.test.ts index 796db91ee..953b94d1b 100644 --- a/js/src/wrappers/groq.test.ts +++ b/js/src/wrappers/groq.test.ts @@ -163,9 +163,8 @@ describe("groq wrapper", () => { }); expect(chatSpan?.metrics).toMatchObject({ completion_tokens: 2, - dram_cached_tokens: 1, + prompt_cached_tokens: 3, prompt_tokens: 5, - sram_cached_tokens: 2, time_to_first_token: expect.any(Number), tokens: 7, }); From d7ae646cd535ea174c10f18fc003aaa5c1a2d558 Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Thu, 20 Aug 2026 14:25:41 +0000 Subject: [PATCH 2/2] Update PR #2391 --- .changeset/groq-batch-api.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .changeset/groq-batch-api.md diff --git a/.changeset/groq-batch-api.md b/.changeset/groq-batch-api.md new file mode 100644 index 000000000..fb7122b69 --- /dev/null +++ b/.changeset/groq-batch-api.md @@ -0,0 +1,5 @@ +--- +"braintrust": minor +--- + +feat: Add OpenAI Batch API Instrumentation