From c05291ffdc8410fe30dc578d611ad5d410b01668 Mon Sep 17 00:00:00 2001 From: Carlos Marchal Date: Fri, 10 Apr 2026 12:36:42 +0200 Subject: [PATCH 1/2] feat: add gemini embedContent tracking support --- .changeset/gemini-embed-content.md | 5 ++ packages/ai/src/gemini/index.ts | 73 ++++++++++++++++ packages/ai/tests/gemini.test.ts | 134 +++++++++++++++++++++++++++++ 3 files changed, 212 insertions(+) create mode 100644 .changeset/gemini-embed-content.md diff --git a/.changeset/gemini-embed-content.md b/.changeset/gemini-embed-content.md new file mode 100644 index 0000000000..b1748cef0d --- /dev/null +++ b/.changeset/gemini-embed-content.md @@ -0,0 +1,5 @@ +--- +'@posthog/ai': minor +--- + +Add Gemini `embedContent` tracking support diff --git a/packages/ai/src/gemini/index.ts b/packages/ai/src/gemini/index.ts index c016ae7fd2..16a1c7fbda 100644 --- a/packages/ai/src/gemini/index.ts +++ b/packages/ai/src/gemini/index.ts @@ -2,12 +2,15 @@ import { GoogleGenAI, GenerateContentResponse as GeminiResponse, GenerateContentParameters, + EmbedContentParameters, + EmbedContentResponse, Part, GenerateContentResponseUsageMetadata, } from '@google/genai' import type { GoogleGenAIOptions } from '@google/genai' import { PostHog } from 'posthog-node' import { + AIEvent, MonitoringParams, sendEventToPosthog, extractAvailableToolCalls, @@ -15,6 +18,7 @@ import { extractPosthogParams, toContentString, sendEventWithErrorToPosthog, + withPrivacyMode, } from '../utils' import { sanitizeGemini } from '../sanitization' import type { TokenUsage, FormattedContent, FormattedContentItem, FormattedMessage } from '../types' @@ -241,6 +245,58 @@ export class WrappedModels { } } + public async embedContent( + params: EmbedContentParameters & MonitoringParams + ): Promise { + const { providerParams: geminiParams, posthogParams } = extractPosthogParams(params) + const startTime = Date.now() + + try { + const response = await this.client.models.embedContent(geminiParams as EmbedContentParameters) + const latency = (Date.now() - startTime) / 1000 + + const inputTokens = extractEmbeddingTokenCount(response) + + await sendEventToPosthog({ + client: this.phClient, + eventType: AIEvent.Embedding, + ...posthogParams, + model: geminiParams.model, + provider: 'gemini', + 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, + }, + }) + + return response + } catch (error: unknown) { + const latency = (Date.now() - startTime) / 1000 + const enrichedError = await sendEventWithErrorToPosthog({ + client: this.phClient, + eventType: AIEvent.Embedding, + ...posthogParams, + model: geminiParams.model, + provider: 'gemini', + input: withPrivacyMode(this.phClient, posthogParams.privacyMode ?? false, geminiParams.contents), + output: null, + latency, + baseURL: 'https://generativelanguage.googleapis.com', + params: params as EmbedContentParameters & MonitoringParams, + usage: { + inputTokens: 0, + }, + error: error, + }) + throw enrichedError + } + } + private formatPartsAsContentBlocks(parts: unknown[]): FormattedContent { const blocks: FormattedContent = [] @@ -383,6 +439,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. diff --git a/packages/ai/tests/gemini.test.ts b/packages/ai/tests/gemini.test.ts index d8b38c2903..84b2241ffb 100644 --- a/packages/ai/tests/gemini.test.ts +++ b/packages/ai/tests/gemini.test.ts @@ -24,6 +24,7 @@ jest.mock('@google/genai', () => { this.models = { generateContent: jest.fn(), generateContentStream: jest.fn(), + embedContent: jest.fn(), } } } @@ -920,4 +921,137 @@ describe('PostHogGemini - Jest test suite', () => { expect(generateContentCall.config.speechConfig.voiceConfig.prebuiltVoiceConfig.voiceName).toBe('Kore') }) }) + + 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 () => { + ;(client as any).client.models.embedContent = jest.fn().mockResolvedValue(mockEmbedResponse) + + const response = await client.models.embedContent({ + model: 'gemini-embedding-001', + contents: 'Hello world', + posthogDistinctId: 'test-id', + posthogProperties: { foo: 'bar' }, + }) + + expect(response).toEqual(mockEmbedResponse) + expect(mockPostHogClient.capture).toHaveBeenCalledTimes(1) + + const [captureArgs] = (mockPostHogClient.capture as jest.Mock).mock.calls + const { distinctId, event, properties } = captureArgs[0] + + expect(distinctId).toBe('test-id') + expect(event).toBe('$ai_embedding') + expect(properties['$ai_provider']).toBe('gemini') + expect(properties['$ai_model']).toBe('gemini-embedding-001') + expect(properties['$ai_input']).toBe('Hello world') + expect(properties['$ai_output_choices']).toBeNull() + expect(properties['$ai_http_status']).toBe(200) + 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('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'], + posthogDistinctId: 'test-id', + }) + + const [captureArgs] = (mockPostHogClient.capture as jest.Mock).mock.calls + expect(captureArgs[0].properties['$ai_input_tokens']).toBe(13) // 5 + 8 + }) + + test('returns 0 tokens when no statistics available', async () => { + ;(client as any).client.models.embedContent = jest.fn().mockResolvedValue(mockEmbedResponse) + + 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(0) + }) + + 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: 'Secret text', + posthogDistinctId: 'test-id', + posthogPrivacyMode: true, + }) + + const [captureArgs] = (mockPostHogClient.capture as jest.Mock).mock.calls + expect(captureArgs[0].properties['$ai_input']).toBeNull() + }) + + 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() + + expect(mockPostHogClient.capture).toHaveBeenCalledTimes(1) + const [captureArgs] = (mockPostHogClient.capture as jest.Mock).mock.calls + 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('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', + posthogDistinctId: 'test-id', + config: { outputDimensionality: 64 }, + }) + + const embedCall = ((client as any).client.models.embedContent as jest.Mock).mock.calls[0][0] + 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']) + }) + }) }) From de75748b3739438472c429503611fecea54ea18d Mon Sep 17 00:00:00 2001 From: Carlos Marchal Date: Fri, 10 Apr 2026 15:01:39 +0200 Subject: [PATCH 2/2] chore: fix prettier formatting --- packages/ai/src/gemini/index.ts | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/packages/ai/src/gemini/index.ts b/packages/ai/src/gemini/index.ts index 16a1c7fbda..d7f16fac19 100644 --- a/packages/ai/src/gemini/index.ts +++ b/packages/ai/src/gemini/index.ts @@ -245,9 +245,7 @@ export class WrappedModels { } } - public async embedContent( - params: EmbedContentParameters & MonitoringParams - ): Promise { + public async embedContent(params: EmbedContentParameters & MonitoringParams): Promise { const { providerParams: geminiParams, posthogParams } = extractPosthogParams(params) const startTime = Date.now()