diff --git a/.changeset/slimy-bars-cheer.md b/.changeset/slimy-bars-cheer.md new file mode 100644 index 000000000..8e552e7bc --- /dev/null +++ b/.changeset/slimy-bars-cheer.md @@ -0,0 +1,5 @@ +--- +"braintrust": minor +--- + +feat: Instrument Google GenAI embedContent for text 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 3916b54df..b94227f03 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 @@ -91,6 +91,39 @@ }, "type": "llm" }, + { + "metadata": { + "operation": "embed" + }, + "metrics": { + "end": 0, + "start": 0 + }, + "name": "google-embed-operation", + "type": null + }, + { + "input": { + "contents": { + "text": "Paris is the capital of France." + }, + "model": "gemini-embedding-001" + }, + "metadata": { + "model": "gemini-embedding-001" + }, + "metrics": { + "duration": 0, + "end": 0, + "start": 0 + }, + "name": "embed_content", + "output": { + "embedding_count": 1, + "embedding_length": 3072 + }, + "type": "llm" + }, { "metadata": { "operation": "attachment" 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 dacc75405..6c19af05e 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 @@ -47,6 +47,38 @@ ], "type": "llm" }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "embed" + }, + "metric_keys": [], + "name": "google-embed-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "gemini-embedding-001" + }, + "metric_keys": [ + "duration" + ], + "name": "embed_content", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, { "has_input": false, "has_output": false, @@ -56,7 +88,7 @@ "metric_keys": [], "name": "google-attachment-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -76,9 +108,9 @@ ], "name": "generate_content", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -91,7 +123,7 @@ "metric_keys": [], "name": "google-stream-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -112,9 +144,9 @@ ], "name": "generate_content_stream", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -127,7 +159,7 @@ "metric_keys": [], "name": "google-stream-return-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -148,9 +180,9 @@ ], "name": "generate_content_stream", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -163,7 +195,7 @@ "metric_keys": [], "name": "google-tool-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -183,9 +215,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 3916b54df..b94227f03 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 @@ -91,6 +91,39 @@ }, "type": "llm" }, + { + "metadata": { + "operation": "embed" + }, + "metrics": { + "end": 0, + "start": 0 + }, + "name": "google-embed-operation", + "type": null + }, + { + "input": { + "contents": { + "text": "Paris is the capital of France." + }, + "model": "gemini-embedding-001" + }, + "metadata": { + "model": "gemini-embedding-001" + }, + "metrics": { + "duration": 0, + "end": 0, + "start": 0 + }, + "name": "embed_content", + "output": { + "embedding_count": 1, + "embedding_length": 3072 + }, + "type": "llm" + }, { "metadata": { "operation": "attachment" 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 dacc75405..6c19af05e 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 @@ -47,6 +47,38 @@ ], "type": "llm" }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "embed" + }, + "metric_keys": [], + "name": "google-embed-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "gemini-embedding-001" + }, + "metric_keys": [ + "duration" + ], + "name": "embed_content", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, { "has_input": false, "has_output": false, @@ -56,7 +88,7 @@ "metric_keys": [], "name": "google-attachment-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -76,9 +108,9 @@ ], "name": "generate_content", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -91,7 +123,7 @@ "metric_keys": [], "name": "google-stream-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -112,9 +144,9 @@ ], "name": "generate_content_stream", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -127,7 +159,7 @@ "metric_keys": [], "name": "google-stream-return-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -148,9 +180,9 @@ ], "name": "generate_content_stream", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -163,7 +195,7 @@ "metric_keys": [], "name": "google-tool-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -183,9 +215,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 3916b54df..b94227f03 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 @@ -91,6 +91,39 @@ }, "type": "llm" }, + { + "metadata": { + "operation": "embed" + }, + "metrics": { + "end": 0, + "start": 0 + }, + "name": "google-embed-operation", + "type": null + }, + { + "input": { + "contents": { + "text": "Paris is the capital of France." + }, + "model": "gemini-embedding-001" + }, + "metadata": { + "model": "gemini-embedding-001" + }, + "metrics": { + "duration": 0, + "end": 0, + "start": 0 + }, + "name": "embed_content", + "output": { + "embedding_count": 1, + "embedding_length": 3072 + }, + "type": "llm" + }, { "metadata": { "operation": "attachment" 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 dacc75405..6c19af05e 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 @@ -47,6 +47,38 @@ ], "type": "llm" }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "embed" + }, + "metric_keys": [], + "name": "google-embed-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "gemini-embedding-001" + }, + "metric_keys": [ + "duration" + ], + "name": "embed_content", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, { "has_input": false, "has_output": false, @@ -56,7 +88,7 @@ "metric_keys": [], "name": "google-attachment-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -76,9 +108,9 @@ ], "name": "generate_content", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -91,7 +123,7 @@ "metric_keys": [], "name": "google-stream-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -112,9 +144,9 @@ ], "name": "generate_content_stream", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -127,7 +159,7 @@ "metric_keys": [], "name": "google-stream-return-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -148,9 +180,9 @@ ], "name": "generate_content_stream", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -163,7 +195,7 @@ "metric_keys": [], "name": "google-tool-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -183,9 +215,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 3916b54df..b94227f03 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 @@ -91,6 +91,39 @@ }, "type": "llm" }, + { + "metadata": { + "operation": "embed" + }, + "metrics": { + "end": 0, + "start": 0 + }, + "name": "google-embed-operation", + "type": null + }, + { + "input": { + "contents": { + "text": "Paris is the capital of France." + }, + "model": "gemini-embedding-001" + }, + "metadata": { + "model": "gemini-embedding-001" + }, + "metrics": { + "duration": 0, + "end": 0, + "start": 0 + }, + "name": "embed_content", + "output": { + "embedding_count": 1, + "embedding_length": 3072 + }, + "type": "llm" + }, { "metadata": { "operation": "attachment" 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 dacc75405..6c19af05e 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 @@ -47,6 +47,38 @@ ], "type": "llm" }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "embed" + }, + "metric_keys": [], + "name": "google-embed-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "gemini-embedding-001" + }, + "metric_keys": [ + "duration" + ], + "name": "embed_content", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, { "has_input": false, "has_output": false, @@ -56,7 +88,7 @@ "metric_keys": [], "name": "google-attachment-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -76,9 +108,9 @@ ], "name": "generate_content", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -91,7 +123,7 @@ "metric_keys": [], "name": "google-stream-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -112,9 +144,9 @@ ], "name": "generate_content_stream", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -127,7 +159,7 @@ "metric_keys": [], "name": "google-stream-return-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -148,9 +180,9 @@ ], "name": "generate_content_stream", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -163,7 +195,7 @@ "metric_keys": [], "name": "google-tool-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -183,9 +215,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 b7c9b47c7..2d66ed8fb 100644 --- a/e2e/scenarios/google-genai-instrumentation/assertions.ts +++ b/e2e/scenarios/google-genai-instrumentation/assertions.ts @@ -13,7 +13,12 @@ import { } from "../../helpers/trace-selectors"; import { summarizeWrapperContract } from "../../helpers/wrapper-contract"; -import { GOOGLE_MODEL, ROOT_NAME, SCENARIO_NAME } from "./scenario.impl.mjs"; +import { + GOOGLE_EMBEDDING_MODEL, + GOOGLE_MODEL, + ROOT_NAME, + SCENARIO_NAME, +} from "./scenario.impl.mjs"; type RunGoogleGenAIScenario = (harness: { runNodeScenarioDir: (options: { @@ -274,6 +279,7 @@ 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 attachmentOperation = findLatestSpan( events, "google-attachment-operation", @@ -292,6 +298,11 @@ function buildRelevantEvents(events: CapturedLogEvent[]): CapturedLogEvent[] { "generate_content", "google-genai.generateContent", ]), + embedOperation, + findGoogleSpan(events, embedOperation?.span.id, [ + "embed_content", + "google-genai.embedContent", + ]), attachmentOperation, findGoogleSpan(events, attachmentOperation?.span.id, [ "generate_content", @@ -391,6 +402,31 @@ export function defineGoogleGenAIInstrumentationAssertions(options: { }, ); + test("captures trace for client.models.embedContent()", testConfig, () => { + const root = findLatestSpan(events, ROOT_NAME); + const operation = findLatestSpan(events, "google-embed-operation"); + const span = findGoogleSpan(events, operation?.span.id, [ + "embed_content", + "google-genai.embedContent", + ]); + + expect(operation).toBeDefined(); + expect(span).toBeDefined(); + expect(operation?.span.parentIds).toEqual([root?.span.id ?? ""]); + expect(span?.row.metadata).toMatchObject({ + model: GOOGLE_EMBEDDING_MODEL, + }); + expect(span?.output).toMatchObject({ + embedding_count: expect.any(Number), + embedding_length: expect.any(Number), + }); + expect(span?.metrics).toMatchObject({ + duration: expect.any(Number), + end: expect.any(Number), + start: 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 716e7dea8..8adf5ca01 100644 --- a/e2e/scenarios/google-genai-instrumentation/scenario.impl.mjs +++ b/e2e/scenarios/google-genai-instrumentation/scenario.impl.mjs @@ -7,6 +7,7 @@ import { } from "../../helpers/provider-runtime.mjs"; const GOOGLE_MODEL = "gemini-2.5-flash-lite"; +const GOOGLE_EMBEDDING_MODEL = "gemini-embedding-001"; const GOOGLE_GROUNDING_MODEL = "gemini-2.0-flash"; const ROOT_NAME = "google-genai-instrumentation-root"; const SCENARIO_NAME = "google-genai-instrumentation"; @@ -138,6 +139,15 @@ async function runGoogleGenAIInstrumentationScenario(sdk, options = {}) { }, GOOGLE_GENAI_RETRY_OPTIONS); }); + await runOperation("google-embed-operation", "embed", async () => { + await withRetry(async () => { + await client.models.embedContent({ + model: GOOGLE_EMBEDDING_MODEL, + contents: "Paris is the capital of France.", + }); + }, GOOGLE_GENAI_RETRY_OPTIONS); + }); + await runOperation( "google-attachment-operation", "attachment", @@ -283,4 +293,4 @@ export async function runAutoGoogleGenAIInstrumentation(sdk) { await runGoogleGenAIInstrumentationScenario(sdk); } -export { GOOGLE_MODEL, ROOT_NAME, SCENARIO_NAME }; +export { GOOGLE_EMBEDDING_MODEL, GOOGLE_MODEL, ROOT_NAME, SCENARIO_NAME }; diff --git a/js/src/auto-instrumentations/configs/google-genai.ts b/js/src/auto-instrumentations/configs/google-genai.ts index e30ba32a3..e1892c35c 100644 --- a/js/src/auto-instrumentations/configs/google-genai.ts +++ b/js/src/auto-instrumentations/configs/google-genai.ts @@ -44,4 +44,35 @@ export const googleGenAIConfigs: InstrumentationConfig[] = [ kind: "Async", }, }, + + // Models.embedContent - class method in older SDK versions + { + channelName: googleGenAIChannels.embedContent.channelName, + module: { + name: "@google/genai", + versionRange: ">=1.0.0 <1.44.0", + filePath: "dist/node/index.mjs", + }, + functionQuery: { + className: "Models", + methodName: "embedContent", + kind: "Async", + }, + }, + + // Models.embedContentInternal - class method in newer SDK versions + // Note: embedContent is an arrow function property that calls this method + { + channelName: googleGenAIChannels.embedContent.channelName, + module: { + name: "@google/genai", + versionRange: ">=1.44.0", + filePath: "dist/node/index.mjs", + }, + functionQuery: { + className: "Models", + methodName: "embedContentInternal", + kind: "Async", + }, + }, ]; diff --git a/js/src/instrumentation/plugins/google-genai-channels.ts b/js/src/instrumentation/plugins/google-genai-channels.ts index 3cf52a482..7d1f9be35 100644 --- a/js/src/instrumentation/plugins/google-genai-channels.ts +++ b/js/src/instrumentation/plugins/google-genai-channels.ts @@ -1,5 +1,7 @@ import { channel, defineChannels } from "../core/channel-definitions"; import type { + GoogleGenAIEmbedContentParams, + GoogleGenAIEmbedContentResponse, GoogleGenAIGenerateContentParams, GoogleGenAIGenerateContentResponse, } from "../../vendor-sdk-types/google-genai"; @@ -25,4 +27,11 @@ export const googleGenAIChannels = defineChannels("@google/genai", { channelName: "models.generateContentStream", kind: "async", }), + embedContent: channel< + [GoogleGenAIEmbedContentParams], + GoogleGenAIEmbedContentResponse + >({ + channelName: "models.embedContent", + kind: "async", + }), }); diff --git a/js/src/instrumentation/plugins/google-genai-plugin.test.ts b/js/src/instrumentation/plugins/google-genai-plugin.test.ts index 62f4f5f2d..8325adca1 100644 --- a/js/src/instrumentation/plugins/google-genai-plugin.test.ts +++ b/js/src/instrumentation/plugins/google-genai-plugin.test.ts @@ -55,18 +55,6 @@ describe("GoogleGenAIPlugin", () => { }); describe("enable/disable lifecycle", () => { - it("should subscribe to channels when enabled", () => { - plugin.enable(); - - expect(mockNewTracingChannel).toHaveBeenCalledWith( - "orchestrion:@google/genai:models.generateContent", - ); - expect(mockNewTracingChannel).toHaveBeenCalledWith( - "orchestrion:@google/genai:models.generateContentStream", - ); - expect(subscribeSpy).toHaveBeenCalled(); - }); - it("should not subscribe multiple times if enabled twice", () => { plugin.enable(); const firstCallCount = subscribeSpy.mock.calls.length; @@ -120,16 +108,6 @@ describe("GoogleGenAIPlugin", () => { expect(handlers).toHaveProperty("error"); }); }); - - describe("generateContentStream channel subscription", () => { - it("should subscribe to streaming channel", () => { - plugin.enable(); - - expect(mockNewTracingChannel).toHaveBeenCalledWith( - "orchestrion:@google/genai:models.generateContentStream", - ); - }); - }); }); describe("Google GenAI serialization functions", () => { diff --git a/js/src/instrumentation/plugins/google-genai-plugin.ts b/js/src/instrumentation/plugins/google-genai-plugin.ts index 09ff9bd40..81c3a6cee 100644 --- a/js/src/instrumentation/plugins/google-genai-plugin.ts +++ b/js/src/instrumentation/plugins/google-genai-plugin.ts @@ -19,6 +19,8 @@ import { SpanTypeAttribute } from "../../../util/index"; import { getCurrentUnixTimestamp } from "../../util"; import { googleGenAIChannels } from "./google-genai-channels"; import type { + GoogleGenAIEmbedContentParams, + GoogleGenAIEmbedContentResponse, GoogleGenAIGenerateContentParams, GoogleGenAIGenerateContentResponse, GoogleGenAIContent, @@ -29,6 +31,10 @@ import type { type GenerateContentChannel = typeof googleGenAIChannels.generateContent; type GenerateContentStreamChannel = typeof googleGenAIChannels.generateContentStream; +type EmbedContentChannel = typeof googleGenAIChannels.embedContent; +type GoogleGenAINonStreamingChannel = + | GenerateContentChannel + | EmbedContentChannel; type GenerateContentStreamEvent = ChannelMessage & { googleGenAIInput?: Record; @@ -65,6 +71,7 @@ function createWrapperParityEvent(args: { * and creates Braintrust spans to track: * - models.generateContent (non-streaming) * - models.generateContentStream (streaming) + * - models.embedContent (embeddings) * * The plugin handles: * - Google-specific token metrics (promptTokenCount, candidatesTokenCount, cachedContentTokenCount) @@ -84,6 +91,7 @@ export class GoogleGenAIPlugin extends BasePlugin { private subscribeToGoogleGenAIChannels(): void { this.subscribeToGenerateContentChannel(); this.subscribeToGenerateContentStreamChannel(); + this.subscribeToEmbedContentChannel(); } private subscribeToGenerateContentChannel(): void { @@ -97,8 +105,8 @@ export class GoogleGenAIPlugin extends BasePlugin { states, (event) => { const params = event.arguments[0]; - const input = serializeInput(params); - const metadata = extractMetadata(params); + const input = serializeGenerateContentInput(params); + const metadata = extractGenerateContentMetadata(params); const span = startSpan({ name: "generate_content", spanAttributes: { @@ -119,8 +127,8 @@ export class GoogleGenAIPlugin extends BasePlugin { start: (event) => { ensureSpanState(states, event, () => { const params = event.arguments[0]; - const input = serializeInput(params); - const metadata = extractMetadata(params); + const input = serializeGenerateContentInput(params); + const metadata = extractGenerateContentMetadata(params); const span = startSpan({ name: "generate_content", spanAttributes: { @@ -182,8 +190,9 @@ export class GoogleGenAIPlugin extends BasePlugin { start: (event) => { const streamEvent = event as GenerateContentStreamEvent; const params = event.arguments[0]; - streamEvent.googleGenAIInput = serializeInput(params); - streamEvent.googleGenAIMetadata = extractMetadata(params); + streamEvent.googleGenAIInput = serializeGenerateContentInput(params); + streamEvent.googleGenAIMetadata = + extractGenerateContentMetadata(params); streamEvent.googleGenAIStartTime = getCurrentUnixTimestamp(); }, asyncEnd: (event) => { @@ -203,6 +212,85 @@ export class GoogleGenAIPlugin extends BasePlugin { tracingChannel.unsubscribe(handlers); }); } + + private subscribeToEmbedContentChannel(): void { + const tracingChannel = + googleGenAIChannels.embedContent.tracingChannel() as IsoTracingChannel< + ChannelMessage + >; + const states = new WeakMap(); + const unbindCurrentSpanStore = bindCurrentSpanStoreToStart( + tracingChannel, + states, + (event) => { + const params = event.arguments[0]; + const input = serializeEmbedContentInput(params); + const metadata = extractEmbedContentMetadata(params); + const span = startSpan({ + name: "embed_content", + spanAttributes: { + type: SpanTypeAttribute.LLM, + }, + event: createWrapperParityEvent({ input, metadata }), + }); + + return { + span, + startTime: getCurrentUnixTimestamp(), + }; + }, + ); + + const handlers: IsoChannelHandlers> = { + start: (event) => { + ensureSpanState(states, event, () => { + const params = event.arguments[0]; + const input = serializeEmbedContentInput(params); + const metadata = extractEmbedContentMetadata(params); + const span = startSpan({ + name: "embed_content", + spanAttributes: { + type: SpanTypeAttribute.LLM, + }, + event: createWrapperParityEvent({ input, metadata }), + }); + + return { + span, + startTime: getCurrentUnixTimestamp(), + }; + }); + }, + asyncEnd: (event) => { + const spanState = states.get(event as object); + if (!spanState) { + return; + } + + try { + const output = summarizeEmbedContentOutput(event.result); + spanState.span.log({ + ...(output ? { output } : {}), + metrics: cleanMetrics( + extractEmbedContentMetrics(event.result, spanState.startTime), + ), + }); + } finally { + spanState.span.end(); + states.delete(event as object); + } + }, + error: (event) => { + logErrorAndEndSpan(states, event as ErrorOf); + }, + }; + + tracingChannel.subscribe(handlers); + this.unsubscribers.push(() => { + unbindCurrentSpanStore?.(); + tracingChannel.unsubscribe(handlers); + }); + } } function ensureSpanState( @@ -220,7 +308,9 @@ function ensureSpanState( return created; } -function bindCurrentSpanStoreToStart( +function bindCurrentSpanStoreToStart< + TChannel extends GoogleGenAINonStreamingChannel, +>( tracingChannel: IsoTracingChannel>, states: WeakMap, create: (event: StartOf) => SpanState, @@ -260,7 +350,7 @@ function bindCurrentSpanStoreToStart( }; } -function logErrorAndEndSpan( +function logErrorAndEndSpan( states: WeakMap, event: ErrorOf, ): void { @@ -497,28 +587,39 @@ function patchGoogleGenAIStreamingResult(args: { return true; } -/** - * Serialize input parameters for Google GenAI API calls. - */ -function serializeInput( +function serializeGenerateContentInput( params: GoogleGenAIGenerateContentParams, ): Record { const input: Record = { model: params.model, - contents: serializeContents(params.contents), + contents: serializeContentCollection(params.contents), }; - if (params.config) { - const config = tryToDict(params.config); - if (config) { - const filteredConfig: Record = {}; - Object.keys(config).forEach((key) => { - if (key !== "tools") { - filteredConfig[key] = config[key]; - } - }); - input.config = filteredConfig; - } + const config = params.config ? tryToDict(params.config) : null; + if (config) { + const filteredConfig: Record = {}; + Object.keys(config).forEach((key) => { + if (key !== "tools") { + filteredConfig[key] = config[key]; + } + }); + input.config = filteredConfig; + } + + return input; +} + +function serializeEmbedContentInput( + params: GoogleGenAIEmbedContentParams, +): Record { + const input: Record = { + model: params.model, + contents: serializeContentCollection(params.contents), + }; + + const config = params.config ? tryToDict(params.config) : null; + if (config) { + input.config = config; } return input; @@ -527,8 +628,8 @@ function serializeInput( /** * Serialize contents, converting inline data to Attachments. */ -function serializeContents( - contents: GoogleGenAIGenerateContentParams["contents"], +function serializeContentCollection( + contents: string | GoogleGenAIContent | GoogleGenAIContent[], ): unknown { if (contents === null || contents === undefined) { return null; @@ -620,32 +721,30 @@ function serializePart(part: GoogleGenAIPart): unknown { return part; } -/** - * Serialize tools configuration. - */ -function serializeTools( +function serializeGenerateContentTools( params: GoogleGenAIGenerateContentParams, ): Record[] | null { - if (!params.config?.tools) { + const config = params.config ? tryToDict(params.config) : null; + const tools = config?.tools; + if (!Array.isArray(tools)) { return null; } try { - return params.config.tools.map((tool) => { - if (typeof tool === "object" && tool.functionDeclarations) { - return tool; + const serializedTools: Record[] = []; + for (const tool of tools) { + const toolDict = tryToDict(tool); + if (toolDict) { + serializedTools.push(toolDict); } - return tool; - }); + } + return serializedTools.length > 0 ? serializedTools : null; } catch { return null; } } -/** - * Extract metadata from parameters. - */ -function extractMetadata( +function extractGenerateContentMetadata( params: GoogleGenAIGenerateContentParams, ): Record { const metadata: Record = {}; @@ -665,7 +764,7 @@ function extractMetadata( } } - const tools = serializeTools(params); + const tools = serializeGenerateContentTools(params); if (tools) { metadata.tools = tools; } @@ -673,6 +772,25 @@ function extractMetadata( return metadata; } +function extractEmbedContentMetadata( + params: GoogleGenAIEmbedContentParams, +): Record { + const metadata: Record = {}; + + if (params.model) { + metadata.model = params.model; + } + + const config = params.config ? tryToDict(params.config) : null; + if (config) { + Object.keys(config).forEach((key) => { + metadata[key] = config[key]; + }); + } + + return metadata; +} + /** * Extract metrics from non-streaming generateContent response. */ @@ -696,6 +814,107 @@ function extractGenerateContentMetrics( return metrics; } +function extractEmbedContentMetrics( + response: GoogleGenAIEmbedContentResponse | undefined, + startTime?: number, +): Record { + const metrics: Record = {}; + + if (startTime !== undefined) { + const end = getCurrentUnixTimestamp(); + metrics.start = startTime; + metrics.end = end; + metrics.duration = end - startTime; + } + + if (response?.usageMetadata) { + populateUsageMetrics(metrics, response.usageMetadata); + } + + const embeddingTokenCount = extractEmbedPromptTokenCount(response); + if (embeddingTokenCount !== undefined) { + metrics.prompt_tokens = embeddingTokenCount; + metrics.tokens = embeddingTokenCount; + } + + return metrics; +} + +function extractEmbedPromptTokenCount( + response: GoogleGenAIEmbedContentResponse | undefined, +): number | undefined { + if (!response) { + return undefined; + } + + // Embedding token counts are available only on Vertex responses via usageMetadata + // and/or embedding.statistics.tokenCount; Gemini Developer API embed responses omit them. + const usagePromptTokens = response.usageMetadata?.promptTokenCount; + if ( + typeof usagePromptTokens === "number" && + Number.isFinite(usagePromptTokens) + ) { + return usagePromptTokens; + } + + const usageTotalTokens = response.usageMetadata?.totalTokenCount; + if ( + typeof usageTotalTokens === "number" && + Number.isFinite(usageTotalTokens) + ) { + return usageTotalTokens; + } + + const embeddings = Array.isArray(response.embeddings) + ? response.embeddings + : response.embedding + ? [response.embedding] + : []; + if (embeddings.length === 0) { + return undefined; + } + + let total = 0; + let sawAny = false; + for (const embedding of embeddings) { + const embeddingStats = tryToDict(tryToDict(embedding)?.statistics); + const tokenCount = embeddingStats?.tokenCount; + if (typeof tokenCount === "number" && Number.isFinite(tokenCount)) { + total += tokenCount; + sawAny = true; + } + } + + return sawAny ? total : undefined; +} + +function summarizeEmbedContentOutput( + response: GoogleGenAIEmbedContentResponse | undefined, +): Record | undefined { + if (!response) { + return undefined; + } + + const embeddings = Array.isArray(response.embeddings) + ? response.embeddings + : response.embedding + ? [response.embedding] + : []; + if (embeddings.length === 0) { + return undefined; + } + + const firstValues = embeddings[0]?.values; + if (!Array.isArray(firstValues)) { + return undefined; + } + + return { + embedding_count: embeddings.length, + embedding_length: firstValues.length, + }; +} + function populateUsageMetrics( metrics: Record, usage: GoogleGenAIUsageMetadata, diff --git a/js/src/vendor-sdk-types/google-genai.ts b/js/src/vendor-sdk-types/google-genai.ts index c0db447c3..d81f1b560 100644 --- a/js/src/vendor-sdk-types/google-genai.ts +++ b/js/src/vendor-sdk-types/google-genai.ts @@ -25,6 +25,9 @@ export interface GoogleGenAIModels { generateContentStream: ( params: GoogleGenAIGenerateContentParams, ) => Promise>; + embedContent: ( + params: GoogleGenAIEmbedContentParams, + ) => Promise; } // Requests @@ -43,6 +46,18 @@ export interface GoogleGenAIGenerateContentParams { [key: string]: unknown; } +export interface GoogleGenAIEmbedContentParams { + model: string; + contents: string | GoogleGenAIContent | GoogleGenAIContent[]; + config?: { + outputDimensionality?: number; + taskType?: string; + toJSON?: () => Record; + [key: string]: unknown; + }; + [key: string]: unknown; +} + export interface GoogleGenAIContent { role?: string; parts: GoogleGenAIPart[]; @@ -78,6 +93,26 @@ export interface GoogleGenAIGenerateContentResponse { [key: string]: unknown; } +export interface GoogleGenAIEmbedding { + values?: number[]; + statistics?: { + tokenCount?: number; + truncated?: boolean; + }; +} + +export interface GoogleGenAIEmbedContentMetadata { + billableCharacterCount?: number; +} + +export interface GoogleGenAIEmbedContentResponse { + embedding?: GoogleGenAIEmbedding; + embeddings?: GoogleGenAIEmbedding[]; + metadata?: GoogleGenAIEmbedContentMetadata; + usageMetadata?: GoogleGenAIUsageMetadata; + [key: string]: unknown; +} + export interface GoogleGenAIGroundingMetadata { groundingChunks?: Array<{ web?: { diff --git a/js/src/wrappers/google-genai.ts b/js/src/wrappers/google-genai.ts index 8a7e5fc61..2a0a4f18c 100644 --- a/js/src/wrappers/google-genai.ts +++ b/js/src/wrappers/google-genai.ts @@ -2,6 +2,7 @@ import { googleGenAIChannels } from "../instrumentation/plugins/google-genai-cha import type { GoogleGenAIClient, GoogleGenAIConstructor, + GoogleGenAIEmbedContentParams, GoogleGenAIGenerateContentParams, GoogleGenAIModels, } from "../vendor-sdk-types/google-genai"; @@ -88,6 +89,8 @@ function wrapModels(models: GoogleGenAIModels): GoogleGenAIModels { return wrapGenerateContentStream( target.generateContentStream.bind(target), ); + } else if (prop === "embedContent") { + return wrapEmbedContent(target.embedContent.bind(target)); } return Reflect.get(target, prop, receiver); }, @@ -117,3 +120,16 @@ function wrapGenerateContentStream( ); }; } + +function wrapEmbedContent( + original: GoogleGenAIModels["embedContent"], +): GoogleGenAIModels["embedContent"] { + return function (params: GoogleGenAIEmbedContentParams) { + return googleGenAIChannels.embedContent.tracePromise( + () => original(params), + { arguments: [params] } as Parameters< + typeof googleGenAIChannels.embedContent.tracePromise + >[1], + ); + }; +}