diff --git a/.changeset/wild-coins-worry.md b/.changeset/wild-coins-worry.md new file mode 100644 index 000000000..ab7f3b344 --- /dev/null +++ b/.changeset/wild-coins-worry.md @@ -0,0 +1,5 @@ +--- +"braintrust": patch +--- + +fix(google-genai): Capture multi-turn message APIs with wrapper diff --git a/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1300.log-payloads.json b/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1300.log-payloads.json index b94227f03..1b803a8a2 100644 --- a/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1300.log-payloads.json +++ b/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1300.log-payloads.json @@ -124,6 +124,165 @@ }, "type": "llm" }, + { + "metadata": { + "operation": "chat" + }, + "metrics": { + "end": 0, + "start": 0 + }, + "name": "google-chat-operation", + "type": null + }, + { + "input": { + "config": { + "maxOutputTokens": 24, + "temperature": 0 + }, + "contents": [ + { + "parts": [ + { + "text": "Reply with exactly MADRID." + } + ], + "role": "user" + } + ], + "model": "gemini-2.5-flash-lite" + }, + "metadata": { + "model": "gemini-2.5-flash-lite" + }, + "metrics": { + "completion_tokens": "", + "duration": 0, + "end": 0, + "prompt_tokens": "", + "start": 0, + "tokens": "" + }, + "name": "generate_content", + "output": { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "MADRID" + } + ], + "role": "model" + }, + "finishReason": "STOP", + "index": 0 + } + ], + "modelVersion": "gemini-2.5-flash-lite", + "responseId": "", + "sdkHttpResponse": { + "headers": { + "alt-svc": "h3=\":443\"; ma=2592000,h3-29=\":443\"; ma=2592000", + "content-encoding": "gzip", + "content-type": "application/json; charset=UTF-8", + "date": "", + "server": "scaffolding on HTTPServer2", + "server-timing": "", + "transfer-encoding": "chunked", + "vary": "Origin, X-Origin, Referer", + "x-content-type-options": "nosniff", + "x-frame-options": "SAMEORIGIN", + "x-gemini-service-tier": "", + "x-xss-protection": "0" + } + }, + "usageMetadata": { + "candidatesTokenCount": "", + "promptTokenCount": "", + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": "" + } + ], + "totalTokenCount": "" + } + }, + "type": "llm" + }, + { + "metadata": { + "operation": "chat-stream" + }, + "metrics": { + "end": 0, + "start": 0 + }, + "name": "google-chat-stream-operation", + "type": null + }, + { + "input": { + "config": { + "maxOutputTokens": 64, + "temperature": 0 + }, + "contents": [ + { + "parts": [ + { + "text": "Count from 1 to 3 and include the words one two three." + } + ], + "role": "user" + } + ], + "model": "gemini-2.5-flash-lite" + }, + "metadata": { + "model": "gemini-2.5-flash-lite" + }, + "metrics": { + "completion_tokens": "", + "duration": 0, + "end": 0, + "prompt_tokens": "", + "start": 0, + "time_to_first_token": 0, + "tokens": "" + }, + "name": "generate_content_stream", + "output": { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "Here's the count from 1 to 3, including the words:\n\n1. One\n2. Two\n3. Three" + } + ], + "role": "model" + }, + "finishReason": "STOP" + } + ], + "text": "Here's the count from 1 to 3, including the words:\n\n1. One\n2. Two\n3. Three", + "usageMetadata": { + "candidatesTokenCount": "", + "promptTokenCount": "", + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": "" + } + ], + "totalTokenCount": "" + } + }, + "type": "llm" + }, { "metadata": { "operation": "attachment" @@ -149,7 +308,7 @@ "url": { "content_type": "image/png", "filename": "file.png", - "key": "", + "key": "", "type": "braintrust_attachment" } } @@ -191,7 +350,7 @@ } ], "modelVersion": "gemini-2.5-flash-lite", - "responseId": "", + "responseId": "", "sdkHttpResponse": { "headers": { "alt-svc": "h3=\":443\"; ma=2592000,h3-29=\":443\"; ma=2592000", @@ -417,7 +576,7 @@ } ], "modelVersion": "gemini-2.5-flash-lite", - "responseId": "", + "responseId": "", "sdkHttpResponse": { "headers": { "alt-svc": "h3=\":443\"; ma=2592000,h3-29=\":443\"; ma=2592000", diff --git a/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1300.span-events.json b/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1300.span-events.json index 6c19af05e..6c7326660 100644 --- a/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1300.span-events.json +++ b/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1300.span-events.json @@ -83,10 +83,10 @@ "has_input": false, "has_output": false, "metadata": { - "operation": "attachment" + "operation": "chat" }, "metric_keys": [], - "name": "google-attachment-operation", + "name": "google-chat-operation", "root_span_id": "", "span_id": "", "span_parents": [ @@ -118,10 +118,10 @@ "has_input": false, "has_output": false, "metadata": { - "operation": "stream" + "operation": "chat-stream" }, "metric_keys": [], - "name": "google-stream-operation", + "name": "google-chat-stream-operation", "root_span_id": "", "span_id": "", "span_parents": [ @@ -150,6 +150,77 @@ ], "type": "llm" }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "attachment" + }, + "metric_keys": [], + "name": "google-attachment-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "gemini-2.5-flash-lite" + }, + "metric_keys": [ + "completion_tokens", + "duration", + "prompt_tokens", + "tokens" + ], + "name": "generate_content", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "stream" + }, + "metric_keys": [], + "name": "google-stream-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "gemini-2.5-flash-lite" + }, + "metric_keys": [ + "completion_tokens", + "duration", + "prompt_tokens", + "time_to_first_token", + "tokens" + ], + "name": "generate_content_stream", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, { "has_input": false, "has_output": false, @@ -159,7 +230,7 @@ "metric_keys": [], "name": "google-stream-return-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -180,9 +251,9 @@ ], "name": "generate_content_stream", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -195,7 +266,7 @@ "metric_keys": [], "name": "google-tool-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -215,9 +286,9 @@ ], "name": "generate_content", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" } diff --git a/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1440.log-payloads.json b/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1440.log-payloads.json index b94227f03..1b803a8a2 100644 --- a/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1440.log-payloads.json +++ b/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1440.log-payloads.json @@ -124,6 +124,165 @@ }, "type": "llm" }, + { + "metadata": { + "operation": "chat" + }, + "metrics": { + "end": 0, + "start": 0 + }, + "name": "google-chat-operation", + "type": null + }, + { + "input": { + "config": { + "maxOutputTokens": 24, + "temperature": 0 + }, + "contents": [ + { + "parts": [ + { + "text": "Reply with exactly MADRID." + } + ], + "role": "user" + } + ], + "model": "gemini-2.5-flash-lite" + }, + "metadata": { + "model": "gemini-2.5-flash-lite" + }, + "metrics": { + "completion_tokens": "", + "duration": 0, + "end": 0, + "prompt_tokens": "", + "start": 0, + "tokens": "" + }, + "name": "generate_content", + "output": { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "MADRID" + } + ], + "role": "model" + }, + "finishReason": "STOP", + "index": 0 + } + ], + "modelVersion": "gemini-2.5-flash-lite", + "responseId": "", + "sdkHttpResponse": { + "headers": { + "alt-svc": "h3=\":443\"; ma=2592000,h3-29=\":443\"; ma=2592000", + "content-encoding": "gzip", + "content-type": "application/json; charset=UTF-8", + "date": "", + "server": "scaffolding on HTTPServer2", + "server-timing": "", + "transfer-encoding": "chunked", + "vary": "Origin, X-Origin, Referer", + "x-content-type-options": "nosniff", + "x-frame-options": "SAMEORIGIN", + "x-gemini-service-tier": "", + "x-xss-protection": "0" + } + }, + "usageMetadata": { + "candidatesTokenCount": "", + "promptTokenCount": "", + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": "" + } + ], + "totalTokenCount": "" + } + }, + "type": "llm" + }, + { + "metadata": { + "operation": "chat-stream" + }, + "metrics": { + "end": 0, + "start": 0 + }, + "name": "google-chat-stream-operation", + "type": null + }, + { + "input": { + "config": { + "maxOutputTokens": 64, + "temperature": 0 + }, + "contents": [ + { + "parts": [ + { + "text": "Count from 1 to 3 and include the words one two three." + } + ], + "role": "user" + } + ], + "model": "gemini-2.5-flash-lite" + }, + "metadata": { + "model": "gemini-2.5-flash-lite" + }, + "metrics": { + "completion_tokens": "", + "duration": 0, + "end": 0, + "prompt_tokens": "", + "start": 0, + "time_to_first_token": 0, + "tokens": "" + }, + "name": "generate_content_stream", + "output": { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "Here's the count from 1 to 3, including the words:\n\n1. One\n2. Two\n3. Three" + } + ], + "role": "model" + }, + "finishReason": "STOP" + } + ], + "text": "Here's the count from 1 to 3, including the words:\n\n1. One\n2. Two\n3. Three", + "usageMetadata": { + "candidatesTokenCount": "", + "promptTokenCount": "", + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": "" + } + ], + "totalTokenCount": "" + } + }, + "type": "llm" + }, { "metadata": { "operation": "attachment" @@ -149,7 +308,7 @@ "url": { "content_type": "image/png", "filename": "file.png", - "key": "", + "key": "", "type": "braintrust_attachment" } } @@ -191,7 +350,7 @@ } ], "modelVersion": "gemini-2.5-flash-lite", - "responseId": "", + "responseId": "", "sdkHttpResponse": { "headers": { "alt-svc": "h3=\":443\"; ma=2592000,h3-29=\":443\"; ma=2592000", @@ -417,7 +576,7 @@ } ], "modelVersion": "gemini-2.5-flash-lite", - "responseId": "", + "responseId": "", "sdkHttpResponse": { "headers": { "alt-svc": "h3=\":443\"; ma=2592000,h3-29=\":443\"; ma=2592000", diff --git a/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1440.span-events.json b/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1440.span-events.json index 6c19af05e..6c7326660 100644 --- a/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1440.span-events.json +++ b/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1440.span-events.json @@ -83,10 +83,10 @@ "has_input": false, "has_output": false, "metadata": { - "operation": "attachment" + "operation": "chat" }, "metric_keys": [], - "name": "google-attachment-operation", + "name": "google-chat-operation", "root_span_id": "", "span_id": "", "span_parents": [ @@ -118,10 +118,10 @@ "has_input": false, "has_output": false, "metadata": { - "operation": "stream" + "operation": "chat-stream" }, "metric_keys": [], - "name": "google-stream-operation", + "name": "google-chat-stream-operation", "root_span_id": "", "span_id": "", "span_parents": [ @@ -150,6 +150,77 @@ ], "type": "llm" }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "attachment" + }, + "metric_keys": [], + "name": "google-attachment-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "gemini-2.5-flash-lite" + }, + "metric_keys": [ + "completion_tokens", + "duration", + "prompt_tokens", + "tokens" + ], + "name": "generate_content", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "stream" + }, + "metric_keys": [], + "name": "google-stream-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "gemini-2.5-flash-lite" + }, + "metric_keys": [ + "completion_tokens", + "duration", + "prompt_tokens", + "time_to_first_token", + "tokens" + ], + "name": "generate_content_stream", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, { "has_input": false, "has_output": false, @@ -159,7 +230,7 @@ "metric_keys": [], "name": "google-stream-return-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -180,9 +251,9 @@ ], "name": "generate_content_stream", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -195,7 +266,7 @@ "metric_keys": [], "name": "google-tool-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -215,9 +286,9 @@ ], "name": "generate_content", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" } diff --git a/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1450.log-payloads.json b/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1450.log-payloads.json index b94227f03..1b803a8a2 100644 --- a/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1450.log-payloads.json +++ b/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1450.log-payloads.json @@ -124,6 +124,165 @@ }, "type": "llm" }, + { + "metadata": { + "operation": "chat" + }, + "metrics": { + "end": 0, + "start": 0 + }, + "name": "google-chat-operation", + "type": null + }, + { + "input": { + "config": { + "maxOutputTokens": 24, + "temperature": 0 + }, + "contents": [ + { + "parts": [ + { + "text": "Reply with exactly MADRID." + } + ], + "role": "user" + } + ], + "model": "gemini-2.5-flash-lite" + }, + "metadata": { + "model": "gemini-2.5-flash-lite" + }, + "metrics": { + "completion_tokens": "", + "duration": 0, + "end": 0, + "prompt_tokens": "", + "start": 0, + "tokens": "" + }, + "name": "generate_content", + "output": { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "MADRID" + } + ], + "role": "model" + }, + "finishReason": "STOP", + "index": 0 + } + ], + "modelVersion": "gemini-2.5-flash-lite", + "responseId": "", + "sdkHttpResponse": { + "headers": { + "alt-svc": "h3=\":443\"; ma=2592000,h3-29=\":443\"; ma=2592000", + "content-encoding": "gzip", + "content-type": "application/json; charset=UTF-8", + "date": "", + "server": "scaffolding on HTTPServer2", + "server-timing": "", + "transfer-encoding": "chunked", + "vary": "Origin, X-Origin, Referer", + "x-content-type-options": "nosniff", + "x-frame-options": "SAMEORIGIN", + "x-gemini-service-tier": "", + "x-xss-protection": "0" + } + }, + "usageMetadata": { + "candidatesTokenCount": "", + "promptTokenCount": "", + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": "" + } + ], + "totalTokenCount": "" + } + }, + "type": "llm" + }, + { + "metadata": { + "operation": "chat-stream" + }, + "metrics": { + "end": 0, + "start": 0 + }, + "name": "google-chat-stream-operation", + "type": null + }, + { + "input": { + "config": { + "maxOutputTokens": 64, + "temperature": 0 + }, + "contents": [ + { + "parts": [ + { + "text": "Count from 1 to 3 and include the words one two three." + } + ], + "role": "user" + } + ], + "model": "gemini-2.5-flash-lite" + }, + "metadata": { + "model": "gemini-2.5-flash-lite" + }, + "metrics": { + "completion_tokens": "", + "duration": 0, + "end": 0, + "prompt_tokens": "", + "start": 0, + "time_to_first_token": 0, + "tokens": "" + }, + "name": "generate_content_stream", + "output": { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "Here's the count from 1 to 3, including the words:\n\n1. One\n2. Two\n3. Three" + } + ], + "role": "model" + }, + "finishReason": "STOP" + } + ], + "text": "Here's the count from 1 to 3, including the words:\n\n1. One\n2. Two\n3. Three", + "usageMetadata": { + "candidatesTokenCount": "", + "promptTokenCount": "", + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": "" + } + ], + "totalTokenCount": "" + } + }, + "type": "llm" + }, { "metadata": { "operation": "attachment" @@ -149,7 +308,7 @@ "url": { "content_type": "image/png", "filename": "file.png", - "key": "", + "key": "", "type": "braintrust_attachment" } } @@ -191,7 +350,7 @@ } ], "modelVersion": "gemini-2.5-flash-lite", - "responseId": "", + "responseId": "", "sdkHttpResponse": { "headers": { "alt-svc": "h3=\":443\"; ma=2592000,h3-29=\":443\"; ma=2592000", @@ -417,7 +576,7 @@ } ], "modelVersion": "gemini-2.5-flash-lite", - "responseId": "", + "responseId": "", "sdkHttpResponse": { "headers": { "alt-svc": "h3=\":443\"; ma=2592000,h3-29=\":443\"; ma=2592000", diff --git a/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1450.span-events.json b/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1450.span-events.json index 6c19af05e..6c7326660 100644 --- a/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1450.span-events.json +++ b/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1450.span-events.json @@ -83,10 +83,10 @@ "has_input": false, "has_output": false, "metadata": { - "operation": "attachment" + "operation": "chat" }, "metric_keys": [], - "name": "google-attachment-operation", + "name": "google-chat-operation", "root_span_id": "", "span_id": "", "span_parents": [ @@ -118,10 +118,10 @@ "has_input": false, "has_output": false, "metadata": { - "operation": "stream" + "operation": "chat-stream" }, "metric_keys": [], - "name": "google-stream-operation", + "name": "google-chat-stream-operation", "root_span_id": "", "span_id": "", "span_parents": [ @@ -150,6 +150,77 @@ ], "type": "llm" }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "attachment" + }, + "metric_keys": [], + "name": "google-attachment-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "gemini-2.5-flash-lite" + }, + "metric_keys": [ + "completion_tokens", + "duration", + "prompt_tokens", + "tokens" + ], + "name": "generate_content", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "stream" + }, + "metric_keys": [], + "name": "google-stream-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "gemini-2.5-flash-lite" + }, + "metric_keys": [ + "completion_tokens", + "duration", + "prompt_tokens", + "time_to_first_token", + "tokens" + ], + "name": "generate_content_stream", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, { "has_input": false, "has_output": false, @@ -159,7 +230,7 @@ "metric_keys": [], "name": "google-stream-return-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -180,9 +251,9 @@ ], "name": "generate_content_stream", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -195,7 +266,7 @@ "metric_keys": [], "name": "google-tool-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -215,9 +286,9 @@ ], "name": "generate_content", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" } diff --git a/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1460.log-payloads.json b/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1460.log-payloads.json index b94227f03..1b803a8a2 100644 --- a/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1460.log-payloads.json +++ b/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1460.log-payloads.json @@ -124,6 +124,165 @@ }, "type": "llm" }, + { + "metadata": { + "operation": "chat" + }, + "metrics": { + "end": 0, + "start": 0 + }, + "name": "google-chat-operation", + "type": null + }, + { + "input": { + "config": { + "maxOutputTokens": 24, + "temperature": 0 + }, + "contents": [ + { + "parts": [ + { + "text": "Reply with exactly MADRID." + } + ], + "role": "user" + } + ], + "model": "gemini-2.5-flash-lite" + }, + "metadata": { + "model": "gemini-2.5-flash-lite" + }, + "metrics": { + "completion_tokens": "", + "duration": 0, + "end": 0, + "prompt_tokens": "", + "start": 0, + "tokens": "" + }, + "name": "generate_content", + "output": { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "MADRID" + } + ], + "role": "model" + }, + "finishReason": "STOP", + "index": 0 + } + ], + "modelVersion": "gemini-2.5-flash-lite", + "responseId": "", + "sdkHttpResponse": { + "headers": { + "alt-svc": "h3=\":443\"; ma=2592000,h3-29=\":443\"; ma=2592000", + "content-encoding": "gzip", + "content-type": "application/json; charset=UTF-8", + "date": "", + "server": "scaffolding on HTTPServer2", + "server-timing": "", + "transfer-encoding": "chunked", + "vary": "Origin, X-Origin, Referer", + "x-content-type-options": "nosniff", + "x-frame-options": "SAMEORIGIN", + "x-gemini-service-tier": "", + "x-xss-protection": "0" + } + }, + "usageMetadata": { + "candidatesTokenCount": "", + "promptTokenCount": "", + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": "" + } + ], + "totalTokenCount": "" + } + }, + "type": "llm" + }, + { + "metadata": { + "operation": "chat-stream" + }, + "metrics": { + "end": 0, + "start": 0 + }, + "name": "google-chat-stream-operation", + "type": null + }, + { + "input": { + "config": { + "maxOutputTokens": 64, + "temperature": 0 + }, + "contents": [ + { + "parts": [ + { + "text": "Count from 1 to 3 and include the words one two three." + } + ], + "role": "user" + } + ], + "model": "gemini-2.5-flash-lite" + }, + "metadata": { + "model": "gemini-2.5-flash-lite" + }, + "metrics": { + "completion_tokens": "", + "duration": 0, + "end": 0, + "prompt_tokens": "", + "start": 0, + "time_to_first_token": 0, + "tokens": "" + }, + "name": "generate_content_stream", + "output": { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "Here's the count from 1 to 3, including the words:\n\n1. One\n2. Two\n3. Three" + } + ], + "role": "model" + }, + "finishReason": "STOP" + } + ], + "text": "Here's the count from 1 to 3, including the words:\n\n1. One\n2. Two\n3. Three", + "usageMetadata": { + "candidatesTokenCount": "", + "promptTokenCount": "", + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": "" + } + ], + "totalTokenCount": "" + } + }, + "type": "llm" + }, { "metadata": { "operation": "attachment" @@ -149,7 +308,7 @@ "url": { "content_type": "image/png", "filename": "file.png", - "key": "", + "key": "", "type": "braintrust_attachment" } } @@ -191,7 +350,7 @@ } ], "modelVersion": "gemini-2.5-flash-lite", - "responseId": "", + "responseId": "", "sdkHttpResponse": { "headers": { "alt-svc": "h3=\":443\"; ma=2592000,h3-29=\":443\"; ma=2592000", @@ -417,7 +576,7 @@ } ], "modelVersion": "gemini-2.5-flash-lite", - "responseId": "", + "responseId": "", "sdkHttpResponse": { "headers": { "alt-svc": "h3=\":443\"; ma=2592000,h3-29=\":443\"; ma=2592000", diff --git a/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1460.span-events.json b/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1460.span-events.json index 6c19af05e..6c7326660 100644 --- a/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1460.span-events.json +++ b/e2e/scenarios/google-genai-instrumentation/__snapshots__/google-genai-v1460.span-events.json @@ -83,10 +83,10 @@ "has_input": false, "has_output": false, "metadata": { - "operation": "attachment" + "operation": "chat" }, "metric_keys": [], - "name": "google-attachment-operation", + "name": "google-chat-operation", "root_span_id": "", "span_id": "", "span_parents": [ @@ -118,10 +118,10 @@ "has_input": false, "has_output": false, "metadata": { - "operation": "stream" + "operation": "chat-stream" }, "metric_keys": [], - "name": "google-stream-operation", + "name": "google-chat-stream-operation", "root_span_id": "", "span_id": "", "span_parents": [ @@ -150,6 +150,77 @@ ], "type": "llm" }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "attachment" + }, + "metric_keys": [], + "name": "google-attachment-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "gemini-2.5-flash-lite" + }, + "metric_keys": [ + "completion_tokens", + "duration", + "prompt_tokens", + "tokens" + ], + "name": "generate_content", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "stream" + }, + "metric_keys": [], + "name": "google-stream-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "gemini-2.5-flash-lite" + }, + "metric_keys": [ + "completion_tokens", + "duration", + "prompt_tokens", + "time_to_first_token", + "tokens" + ], + "name": "generate_content_stream", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, { "has_input": false, "has_output": false, @@ -159,7 +230,7 @@ "metric_keys": [], "name": "google-stream-return-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -180,9 +251,9 @@ ], "name": "generate_content_stream", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -195,7 +266,7 @@ "metric_keys": [], "name": "google-tool-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -215,9 +286,9 @@ ], "name": "generate_content", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" } diff --git a/e2e/scenarios/google-genai-instrumentation/assertions.ts b/e2e/scenarios/google-genai-instrumentation/assertions.ts index 2d66ed8fb..3915de6c6 100644 --- a/e2e/scenarios/google-genai-instrumentation/assertions.ts +++ b/e2e/scenarios/google-genai-instrumentation/assertions.ts @@ -280,6 +280,11 @@ function summarizeGooglePayload(event: CapturedLogEvent): Json { function buildRelevantEvents(events: CapturedLogEvent[]): CapturedLogEvent[] { const generateOperation = findLatestSpan(events, "google-generate-operation"); const embedOperation = findLatestSpan(events, "google-embed-operation"); + const chatOperation = findLatestSpan(events, "google-chat-operation"); + const chatStreamOperation = findLatestSpan( + events, + "google-chat-stream-operation", + ); const attachmentOperation = findLatestSpan( events, "google-attachment-operation", @@ -303,6 +308,16 @@ function buildRelevantEvents(events: CapturedLogEvent[]): CapturedLogEvent[] { "embed_content", "google-genai.embedContent", ]), + chatOperation, + findGoogleSpan(events, chatOperation?.span.id, [ + "generate_content", + "google-genai.generateContent", + ]), + chatStreamOperation, + findGoogleSpan(events, chatStreamOperation?.span.id, [ + "generate_content_stream", + "google-genai.generateContentStream", + ]), attachmentOperation, findGoogleSpan(events, attachmentOperation?.span.id, [ "generate_content", @@ -427,6 +442,48 @@ export function defineGoogleGenAIInstrumentationAssertions(options: { }); }); + test("captures trace for chat.sendMessage()", testConfig, () => { + const root = findLatestSpan(events, ROOT_NAME); + const operation = findLatestSpan(events, "google-chat-operation"); + const span = findGoogleSpan(events, operation?.span.id, [ + "generate_content", + "google-genai.generateContent", + ]); + + expect(operation).toBeDefined(); + expect(span).toBeDefined(); + expect(operation?.span.parentIds).toEqual([root?.span.id ?? ""]); + expect(span?.row.metadata).toMatchObject({ + model: GOOGLE_MODEL, + }); + expect(span?.metrics).toMatchObject({ + duration: expect.any(Number), + end: expect.any(Number), + start: expect.any(Number), + }); + }); + + test("captures trace for chat.sendMessageStream()", testConfig, () => { + const root = findLatestSpan(events, ROOT_NAME); + const operation = findLatestSpan(events, "google-chat-stream-operation"); + const span = findGoogleSpan(events, operation?.span.id, [ + "generate_content_stream", + "google-genai.generateContentStream", + ]); + + expect(operation).toBeDefined(); + expect(span).toBeDefined(); + expect(operation?.span.parentIds).toEqual([root?.span.id ?? ""]); + expect(span?.row.metadata).toMatchObject({ + model: GOOGLE_MODEL, + }); + expect(span?.metrics).toMatchObject({ + time_to_first_token: expect.any(Number), + prompt_tokens: expect.any(Number), + completion_tokens: expect.any(Number), + }); + }); + test("captures trace for sending an attachment", testConfig, () => { const root = findLatestSpan(events, ROOT_NAME); const operation = findLatestSpan(events, "google-attachment-operation"); diff --git a/e2e/scenarios/google-genai-instrumentation/scenario.impl.mjs b/e2e/scenarios/google-genai-instrumentation/scenario.impl.mjs index 8d7a5fdfb..7a35ac72a 100644 --- a/e2e/scenarios/google-genai-instrumentation/scenario.impl.mjs +++ b/e2e/scenarios/google-genai-instrumentation/scenario.impl.mjs @@ -148,6 +148,43 @@ async function runGoogleGenAIInstrumentationScenario(sdk, options = {}) { }, GOOGLE_GENAI_RETRY_OPTIONS); }); + await runOperation("google-chat-operation", "chat", async () => { + await withRetry(async () => { + const chat = client.chats.create({ + model: GOOGLE_MODEL, + config: { + maxOutputTokens: 24, + temperature: 0, + }, + }); + + await chat.sendMessage({ + message: "Reply with exactly MADRID.", + }); + }, GOOGLE_GENAI_RETRY_OPTIONS); + }); + + await runOperation( + "google-chat-stream-operation", + "chat-stream", + async () => { + await withRetry(async () => { + const chat = client.chats.create({ + model: GOOGLE_MODEL, + config: { + maxOutputTokens: 64, + temperature: 0, + }, + }); + + const stream = await chat.sendMessageStream({ + message: "Count from 1 to 3 and include the words one two three.", + }); + await collectAsync(stream); + }, GOOGLE_GENAI_RETRY_OPTIONS); + }, + ); + await runOperation( "google-attachment-operation", "attachment", diff --git a/js/src/vendor-sdk-types/google-genai.ts b/js/src/vendor-sdk-types/google-genai.ts index d81f1b560..b0216c721 100644 --- a/js/src/vendor-sdk-types/google-genai.ts +++ b/js/src/vendor-sdk-types/google-genai.ts @@ -16,6 +16,7 @@ export interface GoogleGenAIConstructor { export interface GoogleGenAIClient { models: GoogleGenAIModels; + chats: GoogleGenAIChats; } export interface GoogleGenAIModels { @@ -30,6 +31,10 @@ export interface GoogleGenAIModels { ) => Promise; } +export interface GoogleGenAIChats { + modelsModule?: GoogleGenAIModels; +} + // Requests export interface GoogleGenAIGenerateContentParams { diff --git a/js/src/wrappers/google-genai.test.ts b/js/src/wrappers/google-genai.test.ts index 8723debfa..4f7b0782a 100644 --- a/js/src/wrappers/google-genai.test.ts +++ b/js/src/wrappers/google-genai.test.ts @@ -273,6 +273,92 @@ describe("google genai client unit tests", TEST_SUITE_OPTIONS, () => { }); }); + test("google genai chat API uses the wrapped models module", async () => { + class FakeChats { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + public constructor(public modelsModule: any) {} + + public create(params: { + model: string; + config?: Record; + }) { + return { + sendMessage: (messageParams: { message: string }) => + this.modelsModule.generateContent({ + model: params.model, + contents: messageParams.message, + config: params.config, + }), + }; + } + } + + class FakeGoogleGenAI { + public models = { + generateContent: async () => ({ + candidates: [ + { + content: { + parts: [{ text: "Hello" }], + role: "model", + }, + finishReason: "STOP", + }, + ], + text: "Hello", + usageMetadata: { + promptTokenCount: 2, + candidatesTokenCount: 1, + totalTokenCount: 3, + }, + }), + }; + + public chats = new FakeChats(this.models); + + public constructor(_config: { apiKey?: string }) {} + } + + initLogger({ + projectName: "google-genai.test.ts", + projectId: "test-project-id", + }); + + const { GoogleGenAI } = wrapGoogleGenAI({ + GoogleGenAI: FakeGoogleGenAI, + }); + const fakeClient = new GoogleGenAI({ apiKey: "test-key" }); + const chat = fakeClient.chats.create({ + model: TEST_MODEL, + config: { + maxOutputTokens: 8, + }, + }); + + const response = await chat.sendMessage({ message: "Say hello." }); + + expect(response.text).toBe("Hello"); + + const spans = await backgroundLogger.drain(); + expect(spans).toHaveLength(1); + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions, @typescript-eslint/no-explicit-any + const span = spans[0] as any; + + expect(span).toMatchObject({ + span_attributes: { + type: "llm", + name: "generate_content", + }, + metadata: expect.objectContaining({ + model: TEST_MODEL, + }), + input: expect.objectContaining({ + model: TEST_MODEL, + contents: expect.anything(), + }), + }); + }); + test("google genai tool calls", async () => { expect(await backgroundLogger.drain()).toHaveLength(0); diff --git a/js/src/wrappers/google-genai.ts b/js/src/wrappers/google-genai.ts index 2a0a4f18c..ecaa23511 100644 --- a/js/src/wrappers/google-genai.ts +++ b/js/src/wrappers/google-genai.ts @@ -1,4 +1,5 @@ import { googleGenAIChannels } from "../instrumentation/plugins/google-genai-channels"; +import { isObject } from "../util"; import type { GoogleGenAIClient, GoogleGenAIConstructor, @@ -70,16 +71,30 @@ function wrapGoogleGenAIClass( function wrapGoogleGenAIInstance( instance: GoogleGenAIClient, ): GoogleGenAIClient { + const wrappedModels = wrapModels(instance.models); + patchGoogleGenAIChats(instance, wrappedModels); + return new Proxy(instance, { get(target, prop, receiver) { if (prop === "models") { - return wrapModels(target.models); + return wrappedModels; } return Reflect.get(target, prop, receiver); }, }); } +function patchGoogleGenAIChats( + instance: GoogleGenAIClient, + wrappedModels: GoogleGenAIModels, +): void { + if (!isObject(instance.chats) || !("modelsModule" in instance.chats)) { + return; + } + + Reflect.set(instance.chats, "modelsModule", wrappedModels); +} + function wrapModels(models: GoogleGenAIModels): GoogleGenAIModels { return new Proxy(models, { get(target, prop, receiver) {