Skip to content

feat: add Gemini embedContent tracking #3369

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
carlos-marchal-ph merged 3 commits into from
Apr 10, 2026
Merged

feat: add Gemini embedContent tracking #3369

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c05291f
feat: add gemini embedContent tracking support
carlos-marchal-ph Apr 10, 2026
de75748
chore: fix prettier formatting
carlos-marchal-ph Apr 10, 2026
181ec80
Merge remote-tracking branch 'origin/main' into feat(llma)/gemini-emb…
carlos-marchal-ph Apr 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .changeset/gemini-embed-content.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
'@posthog/ai': minor
---

Add Gemini `embedContent` tracking support
32 changes: 24 additions & 8 deletions packages/ai/src/gemini/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,22 +2,22 @@ import {
GoogleGenAI,
GenerateContentResponse as GeminiResponse,
GenerateContentParameters,
Part,
GenerateContentResponseUsageMetadata,
EmbedContentParameters,
EmbedContentResponse,
Part,
GenerateContentResponseUsageMetadata,
} from '@google/genai'
import type { GoogleGenAIOptions } from '@google/genai'
import { PostHog } from 'posthog-node'
import {
AIEvent,
MonitoringParams,
sendEventToPosthog,
extractAvailableToolCalls,
formatResponseGemini,
extractPosthogParams,
toContentString,
sendEventWithErrorToPosthog,
AIEvent,
withPrivacyMode,
} from '../utils'
import { sanitizeGemini } from '../sanitization'
Expand Down Expand Up @@ -262,23 +262,22 @@ export class WrappedModels {
const response = await this.client.models.embedContent(geminiParams as EmbedContentParameters)
const latency = (Date.now() - startTime) / 1000

const tokenCount =
response.embeddings?.reduce((sum, embedding) => sum + (embedding.statistics?.tokenCount ?? 0), 0) ?? 0
const inputTokens = extractEmbeddingTokenCount(response)

await sendEventToPosthog({
client: this.phClient,
...posthogParams,
eventType: AIEvent.Embedding,
model: geminiParams.model,
provider: 'gemini',
input: withPrivacyMode(this.phClient, posthogParams.privacyMode, geminiParams.contents),
input: withPrivacyMode(this.phClient, posthogParams.privacyMode ?? false, geminiParams.contents),
output: null,
latency,
baseURL: 'https://generativelanguage.googleapis.com',
params: params as EmbedContentParameters & MonitoringParams,
httpStatus: 200,
usage: {
inputTokens: tokenCount,
inputTokens,
},
})

Expand All @@ -291,7 +290,7 @@ export class WrappedModels {
eventType: AIEvent.Embedding,
model: geminiParams.model,
provider: 'gemini',
input: withPrivacyMode(this.phClient, posthogParams.privacyMode, geminiParams.contents),
input: withPrivacyMode(this.phClient, posthogParams.privacyMode ?? false, geminiParams.contents),
output: null,
latency,
baseURL: 'https://generativelanguage.googleapis.com',
Expand Down Expand Up @@ -447,6 +446,23 @@ export class WrappedModels {
}
}

/**
* Extract total token count from a Gemini embed_content response.
* Token counts are only available per-embedding via Vertex AI's statistics.tokenCount.
* Returns 0 if no token counts are available.
*/
function extractEmbeddingTokenCount(response: EmbedContentResponse): number {
let total = 0
if (response.embeddings) {
for (const embedding of response.embeddings) {
if (embedding.statistics?.tokenCount != null) {
total += embedding.statistics.tokenCount
}
}
}
return total
}

/**
* Detect if Google Search grounding was used in the response.
* Gemini bills per request that uses grounding, not per individual query.
Expand Down
108 changes: 52 additions & 56 deletions packages/ai/tests/gemini.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -923,10 +923,18 @@ describe('PostHogGemini - Jest test suite', () => {
})

describe('embedContent', () => {
const mockEmbedResponse = {
embeddings: [{ values: [0.1, 0.2, 0.3], statistics: null }],
}

const mockEmbedResponseWithStats = {
embeddings: [
{ values: [0.1, 0.2, 0.3], statistics: { tokenCount: 5 } },
{ values: [0.4, 0.5, 0.6], statistics: { tokenCount: 8 } },
],
}

test('basic embedding', async () => {
const mockEmbedResponse = {
embeddings: [{ values: [0.1, 0.2, 0.3] }],
}
;(client as any).client.models.embedContent = jest.fn().mockResolvedValue(mockEmbedResponse)

const response = await client.models.embedContent({
Expand All @@ -949,37 +957,30 @@ describe('PostHogGemini - Jest test suite', () => {
expect(properties['$ai_input']).toBe('Hello world')
expect(properties['$ai_output_choices']).toBeNull()
expect(properties['$ai_http_status']).toBe(200)
expect(properties['foo']).toBe('bar')
expect(typeof properties['$ai_latency']).toBe('number')
expect(properties['$ai_latency']).toBeGreaterThanOrEqual(0)
expect(properties['$ai_trace_id']).toBeDefined()
expect(properties['$ai_base_url']).toBe('https://generativelanguage.googleapis.com')
expect(properties['foo']).toBe('bar')

const embedCall = ((client as any).client.models.embedContent as jest.Mock).mock.calls[0][0]
expect(embedCall.model).toBe('gemini-embedding-001')
expect(embedCall.contents).toBe('Hello world')
})

test('with token counts from statistics', async () => {
const mockEmbedResponse = {
embeddings: [
{ values: [0.1, 0.2], statistics: { tokenCount: 5 } },
{ values: [0.3, 0.4], statistics: { tokenCount: 3 } },
{ values: [0.5, 0.6], statistics: { tokenCount: 7 } },
],
}
;(client as any).client.models.embedContent = jest.fn().mockResolvedValue(mockEmbedResponse)
test('extracts token counts from Vertex AI statistics', async () => {
;(client as any).client.models.embedContent = jest.fn().mockResolvedValue(mockEmbedResponseWithStats)

await client.models.embedContent({
model: 'gemini-embedding-001',
contents: ['Hello', 'World', 'Test'],
contents: ['Hello', 'World'],
posthogDistinctId: 'test-id',
})

const [captureArgs] = (mockPostHogClient.capture as jest.Mock).mock.calls
const { properties } = captureArgs[0]

expect(properties['$ai_input_tokens']).toBe(15)
expect(captureArgs[0].properties['$ai_input_tokens']).toBe(13) // 5 + 8
})

test('without token counts defaults to 0', async () => {
const mockEmbedResponse = {
embeddings: [{ values: [0.1, 0.2, 0.3] }],
}
test('returns 0 tokens when no statistics available', async () => {
;(client as any).client.models.embedContent = jest.fn().mockResolvedValue(mockEmbedResponse)

await client.models.embedContent({
Expand All @@ -989,73 +990,68 @@ describe('PostHogGemini - Jest test suite', () => {
})

const [captureArgs] = (mockPostHogClient.capture as jest.Mock).mock.calls
const { properties } = captureArgs[0]

expect(properties['$ai_input_tokens']).toBe(0)
expect(captureArgs[0].properties['$ai_input_tokens']).toBe(0)
Comment on lines +970 to +993

@greptile-apps greptile-apps Bot Apr 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Prefer parametrised tests for token-count scenarios

These two test cases exercise the same extractEmbeddingTokenCount path with different mock responses and expected values. Per the project's preference, they could be expressed as a single test.each:

test.each([
  ['with Vertex AI statistics', mockEmbedResponseWithStats, 13],
  ['without statistics', mockEmbedResponse, 0],
])('token count %s', async (_, mockResponse, expectedTokens) => {
  ;(client as any).client.models.embedContent = jest.fn().mockResolvedValue(mockResponse)

  await client.models.embedContent({
    model: 'gemini-embedding-001',
    contents: 'Hello',
    posthogDistinctId: 'test-id',
  })

  const [captureArgs] = (mockPostHogClient.capture as jest.Mock).mock.calls
  expect(captureArgs[0].properties['$ai_input_tokens']).toBe(expectedTokens)
})
Prompt To Fix With AI 
This is a comment left during a code review.
Path: packages/ai/tests/gemini.test.ts
Line: 970-993

Comment:
**Prefer parametrised tests for token-count scenarios**

These two test cases exercise the same `extractEmbeddingTokenCount` path with different mock responses and expected values. Per the project's preference, they could be expressed as a single `test.each`:

```ts
test.each([
  ['with Vertex AI statistics', mockEmbedResponseWithStats, 13],
  ['without statistics', mockEmbedResponse, 0],
])('token count %s', async (_, mockResponse, expectedTokens) => {
  ;(client as any).client.models.embedContent = jest.fn().mockResolvedValue(mockResponse)

  await client.models.embedContent({
    model: 'gemini-embedding-001',
    contents: 'Hello',
    posthogDistinctId: 'test-id',
  })

  const [captureArgs] = (mockPostHogClient.capture as jest.Mock).mock.calls
  expect(captureArgs[0].properties['$ai_input_tokens']).toBe(expectedTokens)
})
```

How can I resolve this? If you propose a fix, please make it concise.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

})

test('privacy mode hides input', async () => {
const mockEmbedResponse = {
embeddings: [{ values: [0.1, 0.2] }],
}
test('privacy mode redacts input', async () => {
;(client as any).client.models.embedContent = jest.fn().mockResolvedValue(mockEmbedResponse)

await client.models.embedContent({
model: 'gemini-embedding-001',
contents: 'Sensitive data',
contents: 'Secret text',
posthogDistinctId: 'test-id',
posthogPrivacyMode: true,
})

const [captureArgs] = (mockPostHogClient.capture as jest.Mock).mock.calls
const { properties } = captureArgs[0]

expect(properties['$ai_input']).toBeNull()
expect(properties['$ai_output_choices']).toBeNull()
expect(captureArgs[0].properties['$ai_input']).toBeNull()
})

test('error handling', async () => {
const error = new Error('Embedding API Error')
;(error as any).status = 429
;(client as any).client.models.embedContent = jest.fn().mockRejectedValue(error)
test('error handling captures event and rethrows', async () => {
const mockError = new Error('API error')
;(client as any).client.models.embedContent = jest.fn().mockRejectedValue(mockError)

await expect(
client.models.embedContent({
model: 'gemini-embedding-001',
contents: 'Hello',
posthogDistinctId: 'test-id',
})
).rejects.toThrow('Embedding API Error')
).rejects.toThrow()

expect(mockPostHogClient.capture).toHaveBeenCalledTimes(1)
const [captureArgs] = (mockPostHogClient.capture as jest.Mock).mock.calls
const { event, properties } = captureArgs[0]

expect(event).toBe('$ai_embedding')
expect(properties['$ai_is_error']).toBe(true)
expect(properties['$ai_http_status']).toBe(429)
expect(properties['$ai_input_tokens']).toBe(0)
expect(captureArgs[0].event).toBe('$ai_embedding')
expect(captureArgs[0].properties['$ai_is_error']).toBe(true)
expect(captureArgs[0].properties['$ai_input_tokens']).toBe(0)
})

test('config passed through to underlying call', async () => {
const mockEmbedResponse = {
embeddings: [{ values: [0.1, 0.2] }],
}
test('passes config through to underlying call', async () => {
;(client as any).client.models.embedContent = jest.fn().mockResolvedValue(mockEmbedResponse)

await client.models.embedContent({
model: 'gemini-embedding-001',
contents: 'Hello',
config: { outputDimensionality: 256 },
posthogDistinctId: 'test-id',
} as any)
config: { outputDimensionality: 64 },
})

const embedCall = ((client as any).client.models.embedContent as jest.Mock).mock.calls[0][0]
expect(embedCall.config).toEqual({ outputDimensionality: 256 })
expect(embedCall.model).toBe('gemini-embedding-001')
expect(embedCall.contents).toBe('Hello')
// Posthog params should not be passed through
expect(embedCall.posthogDistinctId).toBeUndefined()
expect(embedCall.config).toEqual({ outputDimensionality: 64 })
})

test('no distinct id sets $process_person_profile to false', async () => {
;(client as any).client.models.embedContent = jest.fn().mockResolvedValue(mockEmbedResponse)

await client.models.embedContent({
model: 'gemini-embedding-001',
contents: 'Hello',
})

const [captureArgs] = (mockPostHogClient.capture as jest.Mock).mock.calls
expect(captureArgs[0].properties['$process_person_profile']).toBe(false)
// distinctId should fall back to traceId
expect(captureArgs[0].distinctId).toBe(captureArgs[0].properties['$ai_trace_id'])
})
})
})
Loading