Crysmags/gemini rebase

.github/workflows/llmobs.yml

@@ -205,3 +205,30 @@ jobs:
         with:
           api_key: ${{ secrets.DD_API_KEY }}
           service: dd-trace-js-tests
+
+  google-genai:
+    runs-on: ubuntu-latest
+    env:
+      PLUGINS: google-genai
+    steps:
+      - uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0
+      - uses: ./.github/actions/testagent/start
+      - uses: ./.github/actions/node/oldest-maintenance-lts
+      - uses: ./.github/actions/install
+      - run: yarn test:plugins:ci
+      - run: yarn test:llmobs:plugins:ci
+        shell: bash
+      - uses: ./.github/actions/node/latest
+      - run: yarn test:plugins:ci
+      - run: yarn test:llmobs:plugins:ci
+        shell: bash
+      - uses: codecov/codecov-action@5a1091511ad55cbe89839c7260b706298ca349f7 # v5.5.1
+      - if: always()
+        uses: ./.github/actions/testagent/logs
+        with:
+          suffix: llmobs-${{ github.job }}
+      - uses: DataDog/junit-upload-github-action@762867566348d59ac9bcf479ebb4ec040db8940a  # v2.0.0
+        if: always()
+        with:
+          api_key: ${{ secrets.DD_API_KEY }}
+          service: dd-trace-js-tests

packages/datadog-instrumentations/src/google-genai.js

@@ -0,0 +1,120 @@
+'use strict'
+
+const { addHook } = require('./helpers/instrument')
+const shimmer = require('../../datadog-shimmer')
+const tracingChannel = require('dc-polyfill').tracingChannel
+const channel = require('dc-polyfill').channel
+
+const genaiTracingChannel = tracingChannel('apm:google:genai:request')
+const onStreamedChunkCh = channel('apm:google:genai:request:chunk')
+
+function wrapGenerateContent (method) {
+  return function wrappedGenerateContent (original) {
+    return function (...args) {
+      if (!genaiTracingChannel.start.hasSubscribers) {
+        return original.apply(this, args)
+      }
+
+      const normalizedName = normalizeMethodName(method)
+
+      const ctx = { args, methodName: normalizedName }
+
+      return genaiTracingChannel.start.runStores(ctx, () => {
+        let result
+        try {
+          result = original.apply(this, arguments)
+        } catch (error) {
+          finish(ctx, null, error)
+          throw error
+        } finally {
+          genaiTracingChannel.end.publish(ctx)
+        }
+        return result.then(response => {
+          if (response[Symbol.asyncIterator]) {
+            shimmer.wrap(response, Symbol.asyncIterator, iterator => wrapStreamIterator(iterator, ctx))
+          } else {
+            finish(ctx, response, null)
+          }
+          return response
+        }).catch(error => {
+          finish(ctx, null, error)
+          throw error
+        })
+      })
+    }
+  }
+}
+
+function wrapStreamIterator (iterator, ctx) {
+  return function () {
+    const itr = iterator.apply(this, arguments)
+    shimmer.wrap(itr, 'next', next => function () {
+      return next.apply(this, arguments)
+        .then(res => {
+          const { done, value: chunk } = res
+          onStreamedChunkCh.publish({ ctx, chunk, done })
+
+          if (done) {
+            finish(ctx)
+          }
+
+          return res
+        })
+        .catch(error => {
+          finish(ctx, null, error)
+          throw error
+        })
+    })
+
+    return itr
+  }
+}
+function finish (ctx, result, error) {
+  if (error) {
+    ctx.error = error
+    genaiTracingChannel.error.publish(ctx)
+  }
+
+  // streamed responses are handled and set separately
+  ctx.result ??= result
+
+  genaiTracingChannel.asyncEnd.publish(ctx)
+}
+// Hook the main package entry point
+addHook({
+  name: '@google/genai',
+  versions: ['>=1.19.0']
+}, exports => {
+  // Wrap GoogleGenAI to intercept when it creates Models instances
+  if (!exports.GoogleGenAI) return exports
+
+  shimmer.wrap(exports, 'GoogleGenAI', GoogleGenAI => {
+    return class extends GoogleGenAI {
+      constructor (...args) {
+        super(...args)
+
+        // We are patching the instance instead of the prototype because when it is compiled from
+        // typescript, the models property is not available on the prototype.
+        if (this.models) {
+          if (this.models.generateContent) {
+            shimmer.wrap(this.models, 'generateContent', wrapGenerateContent('generateContent'))
+          }
+          if (this.models.generateContentStream) {
+            shimmer.wrap(this.models, 'generateContentStream', wrapGenerateContent('generateContentStream'))
+          }
+          if (this.models.embedContent) {
+            shimmer.wrap(this.models, 'embedContent', wrapGenerateContent('embedContent'))
+          }
+        }
+      }
+    }
+  })
+  return exports
+})
+
+function normalizeMethodName (methodName) {
+  // Convert camelCase to snake_case and add Models prefix
+  return 'Models.' + methodName
+    .replaceAll(/([a-z0-9])([A-Z])/g, '$1_$2')
+    .toLowerCase()
+}

packages/datadog-instrumentations/src/helpers/hooks.js

@@ -13,6 +13,7 @@ module.exports = {
   '@playwright/test': () => require('../playwright'),
   '@elastic/elasticsearch': () => require('../elasticsearch'),
   '@elastic/transport': () => require('../elasticsearch'),
+  '@google/genai': () => require('../google-genai'),
   '@google-cloud/pubsub': () => require('../google-cloud-pubsub'),
   '@google-cloud/vertexai': () => require('../google-cloud-vertexai'),
   '@graphql-tools/executor': () => require('../graphql'),

packages/datadog-plugin-google-genai/src/index.js

@@ -0,0 +1,17 @@
+'use strict'
+
+const CompositePlugin = require('../../dd-trace/src/plugins/composite')
+const GenAiTracingPlugin = require('./tracing')
+const GenAiLLMObsPlugin = require('../../dd-trace/src/llmobs/plugins/genai')
+
+class GenAiPlugin extends CompositePlugin {
+  static id = 'google-genai'
+  static get plugins () {
+    return {
+      llmobs: GenAiLLMObsPlugin,
+      tracing: GenAiTracingPlugin
+    }
+  }
+}
+
+module.exports = GenAiPlugin

packages/datadog-plugin-google-genai/src/tracing.js

@@ -0,0 +1,41 @@
+'use strict'
+
+const TracingPlugin = require('../../dd-trace/src/plugins/tracing.js')
+
+class GenAiTracingPlugin extends TracingPlugin {
+  static id = 'google-genai'
+  static operation = 'request'
+  static prefix = 'tracing:apm:google:genai:request'
+
+  static get type () { return 'web' }
+  static get kind () { return 'client' }
+
+  bindStart (ctx) {
+    const { args, methodName } = ctx
+
+    const inputs = args[0]
+    const model = inputs?.model || 'unknown'
+
+    this.startSpan('google_genai.request', {
+      meta: {
+        'resource.name': methodName,
+        'google_genai.request.model': model,
+        'google_genai.request.provider': 'google'
+      }
+    }, ctx)
+
+    return ctx.currentStore
+  }
+
+  asyncEnd (ctx) {
+    const { span } = ctx.currentStore
+    if (!span) return
+
+    if (ctx.result) {
+      span.setTag('google_genai.response.model', ctx.result.modelVersion || ctx.inputs?.model)
+    }
+    span.finish()
+  }
+}
+
+module.exports = GenAiTracingPlugin

packages/datadog-plugin-google-genai/test/index.spec.js

@@ -0,0 +1,91 @@
+'use strict'
+
+const { describe, before, after, it } = require('mocha')
+const { withVersions } = require('../../dd-trace/test/setup/mocha')
+const agent = require('../../dd-trace/test/plugins/agent')
+const assert = require('node:assert')
+
+describe('Plugin', () => {
+  withVersions('google-genai', '@google/genai', (version) => {
+    let client
+
+    before(async () => {
+      await agent.load('google-genai')
+
+      const { GoogleGenAI } = require(`../../../versions/@google/genai@${version}`).get()
+      client = new GoogleGenAI({
+        apiKey: process.env.GOOGLE_API_KEY || '<not-a-real-key>',
+        httpOptions: { baseUrl: 'http://127.0.0.1:9126/vcr/genai' }
+      })
+    })
+
+    after(async () => {
+      await agent.close({ ritmReset: false })
+    })
+
+    describe('models.generateContent', () => {
+      it('creates a span', async () => {
+        const tracesPromise = agent.assertSomeTraces(traces => {
+          const span = traces[0][0]
+
+          assert.equal(span.name, 'google_genai.request')
+          assert.equal(span.resource, 'Models.generate_content')
+          assert.equal(span.meta['google_genai.request.model'], 'gemini-2.0-flash')
+        })
+
+        const result = await client.models.generateContent({
+          model: 'gemini-2.0-flash',
+          contents: 'Hello, world!'
+        })
+
+        assert.ok(result)
+
+        await tracesPromise
+      })
+    })
+
+    describe('models.generateContentStream', () => {
+      it('creates a span', async () => {
+        const tracesPromise = agent.assertSomeTraces(traces => {
+          const span = traces[0][0]
+
+          assert.equal(span.name, 'google_genai.request')
+          assert.equal(span.resource, 'Models.generate_content_stream')
+          assert.equal(span.meta['google_genai.request.model'], 'gemini-2.0-flash')
+        })
+
+        const stream = await client.models.generateContentStream({
+          model: 'gemini-2.0-flash',
+          contents: 'Hello, world!'
+        })
+
+        for await (const chunk of stream) {
+          assert.ok(chunk)
+        }
+
+        await tracesPromise
+      })
+    })
+
+    describe('models.embedContent', () => {
+      it('creates a span', async () => {
+        const tracesPromise = agent.assertSomeTraces(traces => {
+          const span = traces[0][0]
+
+          assert.equal(span.name, 'google_genai.request')
+          assert.equal(span.resource, 'Models.embed_content')
+          assert.equal(span.meta['google_genai.request.model'], 'text-embedding-004')
+        })
+
+        const result = await client.models.embedContent({
+          model: 'text-embedding-004',
+          contents: 'Hello, world!'
+        })
+
+        assert.ok(result)
+
+        await tracesPromise
+      })
+    })
+  })
+})

packages/datadog-plugin-google-genai/test/integration-test/client.spec.js

@@ -0,0 +1,49 @@
+'use strict'
+
+const {
+  FakeAgent,
+  sandboxCwd,
+  useSandbox,
+  checkSpansForServiceName,
+  spawnPluginIntegrationTestProc
+} = require('../../../../integration-tests/helpers')
+const { withVersions } = require('../../../dd-trace/test/setup/mocha')
+const { assert } = require('chai')
+const { describe, it, beforeEach, afterEach } = require('mocha')
+
+describe('esm', () => {
+  let agent
+  let proc
+
+  withVersions('google-genai', ['@google/genai'], version => {
+    useSandbox([
+      `@google/genai@${version}`,
+    ], false, [
+      './packages/datadog-plugin-google-genai/test/integration-test/*'
+    ])
+
+    beforeEach(async () => {
+      agent = await new FakeAgent().start()
+    })
+
+    afterEach(async () => {
+      proc?.kill()
+      await agent.stop()
+    })
+
+    it('is instrumented', async () => {
+      const res = agent.assertMessageReceived(({ headers, payload }) => {
+        assert.propertyVal(headers, 'host', `127.0.0.1:${agent.port}`)
+        assert.isArray(payload)
+        assert.strictEqual(checkSpansForServiceName(payload, 'google_genai.request'), true)
+      })
+
+      proc = await spawnPluginIntegrationTestProc(sandboxCwd(), 'server.mjs', agent.port, null, {
+        NODE_OPTIONS: '--import dd-trace/initialize.mjs',
+        GOOGLE_API_KEY: process.env.GOOGLE_API_KEY || '<not-a-real-key>'
+      })
+
+      await res
+    }).timeout(20000)
+  })
+})

packages/datadog-plugin-google-genai/test/integration-test/server.mjs

@@ -0,0 +1,11 @@
+import { GoogleGenAI } from '@google/genai'
+
+const client = new GoogleGenAI({
+  apiKey: process.env.GOOGLE_API_KEY,
+  httpOptions: { baseUrl: 'http://127.0.0.1:9126/vcr/genai' }
+})
+
+await client.models.generateContent({
+  model: 'gemini-2.0-flash',
+  contents: 'Hello, world!'
+})