@posthog/ai 4.4.0 → 5.0.1

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@posthog/ai",
-  "version": "4.4.0",
+  "version": "5.0.1",
   "description": "PostHog Node.js AI integrations",
   "repository": {
     "type": "git",
@@ -13,7 +13,7 @@
   "license": "MIT",
   "devDependencies": {
     "@types/jest": "^29.5.14",
-    "@types/node": "^18.0.0",
+    "@types/node": "^20.0.0",
     "node-fetch": "^3.3.2",
     "typescript": "^4.7.4"
   },
@@ -22,11 +22,16 @@
     "ai",
     "openai",
     "anthropic",
+    "gemini",
     "llm",
     "observability"
   ],
+  "engines": {
+    "node": ">=20"
+  },
   "dependencies": {
     "@anthropic-ai/sdk": "^0.36.3",
+    "@google/genai": "^1.1.0",
     "@langchain/core": "^0.3.37",
     "ai": "^4.1.0",
     "langchain": "^0.3.15",
@@ -35,7 +40,7 @@
     "zod": "^3.24.1"
   },
   "peerDependencies": {
-    "posthog-node": "^4.17.1"
+    "posthog-node": "^5.0.0"
   },
   "scripts": {
     "test": "jest",
@@ -52,6 +57,11 @@
       "import": "./lib/anthropic/index.mjs",
       "types": "./lib/anthropic/index.d.ts"
     },
+    "./gemini": {
+      "require": "./lib/gemini/index.cjs",
+      "import": "./lib/gemini/index.mjs",
+      "types": "./lib/gemini/index.d.ts"
+    },
     "./openai": {
       "require": "./lib/openai/index.cjs",
       "import": "./lib/openai/index.mjs",

package/src/gemini/index.ts

@@ -0,0 +1,254 @@
+import { GoogleGenAI } from '@google/genai'
+import { PostHog } from 'posthog-node'
+import { v4 as uuidv4 } from 'uuid'
+import { MonitoringParams, sendEventToPosthog } from '../utils'
+
+// Types from @google/genai
+type GenerateContentRequest = {
+  model: string
+  contents: any
+  config?: any
+  [key: string]: any
+}
+
+type GenerateContentResponse = {
+  text?: string
+  candidates?: any[]
+  usageMetadata?: {
+    promptTokenCount?: number
+    candidatesTokenCount?: number
+    totalTokenCount?: number
+  }
+  [key: string]: any
+}
+
+interface MonitoringGeminiConfig {
+  apiKey?: string
+  vertexai?: boolean
+  project?: string
+  location?: string
+  apiVersion?: string
+  posthog: PostHog
+}
+
+export class PostHogGoogleGenAI {
+  private readonly phClient: PostHog
+  private readonly client: GoogleGenAI
+  public models: WrappedModels
+
+  constructor(config: MonitoringGeminiConfig) {
+    const { posthog, ...geminiConfig } = config
+    this.phClient = posthog
+    this.client = new GoogleGenAI(geminiConfig)
+    this.models = new WrappedModels(this.client, this.phClient)
+  }
+}
+
+export class WrappedModels {
+  private readonly phClient: PostHog
+  private readonly client: GoogleGenAI
+
+  constructor(client: GoogleGenAI, phClient: PostHog) {
+    this.client = client
+    this.phClient = phClient
+  }
+
+  public async generateContent(params: GenerateContentRequest & MonitoringParams): Promise<GenerateContentResponse> {
+    const {
+      posthogDistinctId,
+      posthogTraceId,
+      posthogProperties,
+      posthogGroups,
+      posthogCaptureImmediate,
+      ...geminiParams
+    } = params
+
+    const traceId = posthogTraceId ?? uuidv4()
+    const startTime = Date.now()
+
+    try {
+      const response = await this.client.models.generateContent(geminiParams)
+      const latency = (Date.now() - startTime) / 1000
+
+      await sendEventToPosthog({
+        client: this.phClient,
+        distinctId: posthogDistinctId ?? traceId,
+        traceId,
+        model: geminiParams.model,
+        provider: 'gemini',
+        input: this.formatInput(geminiParams.contents),
+        output: this.formatOutput(response),
+        latency,
+        baseURL: 'https://generativelanguage.googleapis.com',
+        params: params as any,
+        httpStatus: 200,
+        usage: {
+          inputTokens: response.usageMetadata?.promptTokenCount ?? 0,
+          outputTokens: response.usageMetadata?.candidatesTokenCount ?? 0,
+        },
+        captureImmediate: posthogCaptureImmediate,
+      })
+
+      return response
+    } catch (error: any) {
+      const latency = (Date.now() - startTime) / 1000
+      await sendEventToPosthog({
+        client: this.phClient,
+        distinctId: posthogDistinctId ?? traceId,
+        traceId,
+        model: geminiParams.model,
+        provider: 'gemini',
+        input: this.formatInput(geminiParams.contents),
+        output: [],
+        latency,
+        baseURL: 'https://generativelanguage.googleapis.com',
+        params: params as any,
+        httpStatus: error?.status ?? 500,
+        usage: {
+          inputTokens: 0,
+          outputTokens: 0,
+        },
+        isError: true,
+        error: JSON.stringify(error),
+        captureImmediate: posthogCaptureImmediate,
+      })
+      throw error
+    }
+  }
+
+  public async *generateContentStream(
+    params: GenerateContentRequest & MonitoringParams
+  ): AsyncGenerator<any, void, unknown> {
+    const {
+      posthogDistinctId,
+      posthogTraceId,
+      posthogProperties,
+      posthogGroups,
+      posthogCaptureImmediate,
+      ...geminiParams
+    } = params
+
+    const traceId = posthogTraceId ?? uuidv4()
+    const startTime = Date.now()
+    let accumulatedContent = ''
+    let usage = {
+      inputTokens: 0,
+      outputTokens: 0,
+    }
+
+    try {
+      const stream = await this.client.models.generateContentStream(geminiParams)
+
+      for await (const chunk of stream) {
+        if (chunk.text) {
+          accumulatedContent += chunk.text
+        }
+        if (chunk.usageMetadata) {
+          usage = {
+            inputTokens: chunk.usageMetadata.promptTokenCount ?? 0,
+            outputTokens: chunk.usageMetadata.candidatesTokenCount ?? 0,
+          }
+        }
+        yield chunk
+      }
+
+      const latency = (Date.now() - startTime) / 1000
+      await sendEventToPosthog({
+        client: this.phClient,
+        distinctId: posthogDistinctId ?? traceId,
+        traceId,
+        model: geminiParams.model,
+        provider: 'gemini',
+        input: this.formatInput(geminiParams.contents),
+        output: [{ content: accumulatedContent, role: 'assistant' }],
+        latency,
+        baseURL: 'https://generativelanguage.googleapis.com',
+        params: params as any,
+        httpStatus: 200,
+        usage,
+        captureImmediate: posthogCaptureImmediate,
+      })
+    } catch (error: any) {
+      const latency = (Date.now() - startTime) / 1000
+      await sendEventToPosthog({
+        client: this.phClient,
+        distinctId: posthogDistinctId ?? traceId,
+        traceId,
+        model: geminiParams.model,
+        provider: 'gemini',
+        input: this.formatInput(geminiParams.contents),
+        output: [],
+        latency,
+        baseURL: 'https://generativelanguage.googleapis.com',
+        params: params as any,
+        httpStatus: error?.status ?? 500,
+        usage: {
+          inputTokens: 0,
+          outputTokens: 0,
+        },
+        isError: true,
+        error: JSON.stringify(error),
+        captureImmediate: posthogCaptureImmediate,
+      })
+      throw error
+    }
+  }
+
+  private formatInput(contents: any): Array<{ role: string; content: string }> {
+    if (typeof contents === 'string') {
+      return [{ role: 'user', content: contents }]
+    }
+
+    if (Array.isArray(contents)) {
+      return contents.map((item) => {
+        if (typeof item === 'string') {
+          return { role: 'user', content: item }
+        }
+        if (item && typeof item === 'object') {
+          if (item.text) {
+            return { role: item.role || 'user', content: item.text }
+          }
+          if (item.content) {
+            return { role: item.role || 'user', content: item.content }
+          }
+        }
+        return { role: 'user', content: String(item) }
+      })
+    }
+
+    if (contents && typeof contents === 'object') {
+      if (contents.text) {
+        return [{ role: 'user', content: contents.text }]
+      }
+      if (contents.content) {
+        return [{ role: 'user', content: contents.content }]
+      }
+    }
+
+    return [{ role: 'user', content: String(contents) }]
+  }
+
+  private formatOutput(response: GenerateContentResponse): Array<{ role: string; content: string }> {
+    if (response.text) {
+      return [{ role: 'assistant', content: response.text }]
+    }
+
+    if (response.candidates && Array.isArray(response.candidates)) {
+      return response.candidates.map((candidate) => {
+        if (candidate.content && candidate.content.parts) {
+          const text = candidate.content.parts
+            .filter((part: any) => part.text)
+            .map((part: any) => part.text)
+            .join('')
+          return { role: 'assistant', content: text }
+        }
+        return { role: 'assistant', content: String(candidate) }
+      })
+    }
+
+    return []
+  }
+}
+
+export default PostHogGoogleGenAI
+export { PostHogGoogleGenAI as Gemini }

package/src/index.ts

@@ -2,10 +2,12 @@ import PostHogOpenAI from './openai'
 import PostHogAzureOpenAI from './openai/azure'
 import { wrapVercelLanguageModel } from './vercel/middleware'
 import PostHogAnthropic from './anthropic'
+import PostHogGoogleGenAI from './gemini'
 import { LangChainCallbackHandler } from './langchain/callbacks'
 
 export { PostHogOpenAI as OpenAI }
 export { PostHogAzureOpenAI as AzureOpenAI }
 export { PostHogAnthropic as Anthropic }
+export { PostHogGoogleGenAI as GoogleGenAI }
 export { wrapVercelLanguageModel as withTracing }
 export { LangChainCallbackHandler }

package/src/utils.ts

@@ -66,6 +66,8 @@ export const formatResponse = (response: any, provider: string): Array<{ role: s
     return formatResponseAnthropic(response)
   } else if (provider === 'openai') {
     return formatResponseOpenAI(response)
+  } else if (provider === 'gemini') {
+    return formatResponseGemini(response)
   }
   return []
 }
@@ -97,6 +99,37 @@ export const formatResponseOpenAI = (response: any): Array<{ role: string; conte
   return output
 }
 
+export const formatResponseGemini = (response: any): Array<{ role: string; content: string }> => {
+  const output: Array<{ role: string; content: string }> = []
+
+  if (response.text) {
+    output.push({
+      role: 'assistant',
+      content: response.text,
+    })
+    return output
+  }
+
+  if (response.candidates && Array.isArray(response.candidates)) {
+    for (const candidate of response.candidates) {
+      if (candidate.content && candidate.content.parts) {
+        const text = candidate.content.parts
+          .filter((part: any) => part.text)
+          .map((part: any) => part.text)
+          .join('')
+        if (text) {
+          output.push({
+            role: 'assistant',
+            content: text,
+          })
+        }
+      }
+    }
+  }
+
+  return output
+}
+
 export const mergeSystemPrompt = (params: MessageCreateParams & MonitoringParams, provider: string): any => {
   if (provider == 'anthropic') {
     const messages = params.messages || []

package/tests/gemini.test.ts

@@ -0,0 +1,313 @@
+import { PostHog } from 'posthog-node'
+import PostHogGemini from '../src/gemini'
+
+let mockGeminiResponse: any = {}
+let mockGeminiStreamResponse: any = {}
+
+jest.mock('posthog-node', () => {
+  return {
+    PostHog: jest.fn().mockImplementation(() => {
+      return {
+        capture: jest.fn(),
+        captureImmediate: jest.fn(),
+        privacyMode: false,
+      }
+    }),
+  }
+})
+
+jest.mock('@google/genai', () => {
+  class MockGoogleGenAI {
+    models: any
+    constructor() {
+      this.models = {
+        generateContent: jest.fn(),
+        generateContentStream: jest.fn(),
+      }
+    }
+  }
+
+  return {
+    GoogleGenAI: MockGoogleGenAI,
+  }
+})
+
+describe('PostHogGemini - Jest test suite', () => {
+  let mockPostHogClient: PostHog
+  let client: PostHogGemini
+
+  beforeAll(() => {
+    if (!process.env.GEMINI_API_KEY) {
+      console.warn('⚠️ Skipping Gemini tests: No GEMINI_API_KEY environment variable set')
+    }
+  })
+
+  beforeEach(() => {
+    // Skip all tests if no API key is present
+    if (!process.env.GEMINI_API_KEY) {
+      return
+    }
+
+    jest.clearAllMocks()
+
+    // Reset the default mocks
+    mockPostHogClient = new (PostHog as any)()
+    client = new PostHogGemini({
+      apiKey: process.env.GEMINI_API_KEY || '',
+      posthog: mockPostHogClient as any,
+    })
+
+    // Some default generate content mock
+    mockGeminiResponse = {
+      text: 'Hello from Gemini!',
+      candidates: [
+        {
+          content: {
+            parts: [
+              {
+                text: 'Hello from Gemini!',
+              },
+            ],
+          },
+          finishReason: 'STOP',
+        },
+      ],
+      usageMetadata: {
+        promptTokenCount: 15,
+        candidatesTokenCount: 8,
+        totalTokenCount: 23,
+      },
+    }
+
+    // Mock streaming response
+    mockGeminiStreamResponse = [
+      {
+        text: 'Hello ',
+        candidates: [
+          {
+            content: {
+              parts: [{ text: 'Hello ' }],
+            },
+          },
+        ],
+        usageMetadata: {
+          promptTokenCount: 15,
+          candidatesTokenCount: 2,
+        },
+      },
+      {
+        text: 'from ',
+        candidates: [
+          {
+            content: {
+              parts: [{ text: 'from ' }],
+            },
+          },
+        ],
+        usageMetadata: {
+          promptTokenCount: 15,
+          candidatesTokenCount: 4,
+        },
+      },
+      {
+        text: 'Gemini!',
+        candidates: [
+          {
+            content: {
+              parts: [{ text: 'Gemini!' }],
+            },
+          },
+        ],
+        usageMetadata: {
+          promptTokenCount: 15,
+          candidatesTokenCount: 8,
+        },
+      },
+    ]
+
+    // Mock the generateContent method
+    ;(client as any).client.models.generateContent = jest.fn().mockResolvedValue(mockGeminiResponse)
+
+    // Mock the generateContentStream method
+    ;(client as any).client.models.generateContentStream = jest.fn().mockImplementation(async function* () {
+      for (const chunk of mockGeminiStreamResponse) {
+        yield chunk
+      }
+    })
+  })
+
+  // Wrap each test with conditional skip
+  const conditionalTest = process.env.GEMINI_API_KEY ? test : test.skip
+
+  conditionalTest('basic content generation', async () => {
+    const response = await client.models.generateContent({
+      model: 'gemini-2.0-flash-001',
+      contents: 'Hello',
+      posthogDistinctId: 'test-id',
+      posthogProperties: { foo: 'bar' },
+    })
+
+    expect(response).toEqual(mockGeminiResponse)
+    // We expect 1 capture call
+    expect(mockPostHogClient.capture).toHaveBeenCalledTimes(1)
+    // Check the capture arguments
+    const [captureArgs] = (mockPostHogClient.capture as jest.Mock).mock.calls
+    const { distinctId, event, properties } = captureArgs[0]
+
+    expect(distinctId).toBe('test-id')
+    expect(event).toBe('$ai_generation')
+    expect(properties['$ai_provider']).toBe('gemini')
+    expect(properties['$ai_model']).toBe('gemini-2.0-flash-001')
+    expect(properties['$ai_input']).toEqual([{ role: 'user', content: 'Hello' }])
+    expect(properties['$ai_output_choices']).toEqual([{ role: 'assistant', content: 'Hello from Gemini!' }])
+    expect(properties['$ai_input_tokens']).toBe(15)
+    expect(properties['$ai_output_tokens']).toBe(8)
+    expect(properties['$ai_http_status']).toBe(200)
+    expect(properties['foo']).toBe('bar')
+    expect(typeof properties['$ai_latency']).toBe('number')
+  })
+
+  conditionalTest('streaming content generation', async () => {
+    const stream = client.models.generateContentStream({
+      model: 'gemini-2.0-flash-001',
+      contents: 'Write a short poem',
+      posthogDistinctId: 'test-id',
+      posthogProperties: { foo: 'bar' },
+    })
+
+    let accumulatedText = ''
+    for await (const chunk of stream) {
+      if (chunk.text) {
+        accumulatedText += chunk.text
+      }
+    }
+
+    expect(accumulatedText).toBe('Hello from Gemini!')
+    // We expect 1 capture call after streaming completes
+    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_generation')
+    expect(properties['$ai_provider']).toBe('gemini')
+    expect(properties['$ai_model']).toBe('gemini-2.0-flash-001')
+    expect(properties['$ai_input']).toEqual([{ role: 'user', content: 'Write a short poem' }])
+    expect(properties['$ai_output_choices']).toEqual([{ content: 'Hello from Gemini!', role: 'assistant' }])
+    expect(properties['$ai_input_tokens']).toBe(15)
+    expect(properties['$ai_output_tokens']).toBe(8)
+    expect(properties['$ai_http_status']).toBe(200)
+    expect(properties['foo']).toBe('bar')
+    expect(typeof properties['$ai_latency']).toBe('number')
+  })
+
+  conditionalTest('groups', async () => {
+    await client.models.generateContent({
+      model: 'gemini-2.0-flash-001',
+      contents: 'Hello',
+      posthogDistinctId: 'test-id',
+      posthogGroups: { team: 'ai-team' },
+    })
+
+    expect(mockPostHogClient.capture).toHaveBeenCalledTimes(1)
+    const [captureArgs] = (mockPostHogClient.capture as jest.Mock).mock.calls
+    const { groups } = captureArgs[0]
+
+    expect(groups).toEqual({ team: 'ai-team' })
+  })
+
+  conditionalTest('privacy mode', async () => {
+    await client.models.generateContent({
+      model: 'gemini-2.0-flash-001',
+      contents: 'Sensitive information',
+      posthogDistinctId: 'test-id',
+      posthogPrivacyMode: true,
+    })
+
+    expect(mockPostHogClient.capture).toHaveBeenCalledTimes(1)
+    const [captureArgs] = (mockPostHogClient.capture as jest.Mock).mock.calls
+    const { properties } = captureArgs[0]
+
+    expect(properties['$ai_input']).toBeNull()
+    expect(properties['$ai_output_choices']).toBeNull()
+  })
+
+  conditionalTest('error handling', async () => {
+    const error = new Error('API Error')
+    ;(error as any).status = 400
+    ;(client as any).client.models.generateContent = jest.fn().mockRejectedValue(error)
+
+    await expect(
+      client.models.generateContent({
+        model: 'gemini-2.0-flash-001',
+        contents: 'Hello',
+        posthogDistinctId: 'test-id',
+      })
+    ).rejects.toThrow('API Error')
+
+    expect(mockPostHogClient.capture).toHaveBeenCalledTimes(1)
+    const [captureArgs] = (mockPostHogClient.capture as jest.Mock).mock.calls
+    const { properties } = captureArgs[0]
+
+    expect(properties['$ai_is_error']).toBe(true)
+    expect(properties['$ai_http_status']).toBe(400)
+    expect(properties['$ai_input_tokens']).toBe(0)
+    expect(properties['$ai_output_tokens']).toBe(0)
+  })
+
+  conditionalTest('array contents input', async () => {
+    await client.models.generateContent({
+      model: 'gemini-2.0-flash-001',
+      contents: ['Hello', 'How are you?'],
+      posthogDistinctId: 'test-id',
+    })
+
+    expect(mockPostHogClient.capture).toHaveBeenCalledTimes(1)
+    const [captureArgs] = (mockPostHogClient.capture as jest.Mock).mock.calls
+    const { properties } = captureArgs[0]
+
+    expect(properties['$ai_input']).toEqual([
+      { role: 'user', content: 'Hello' },
+      { role: 'user', content: 'How are you?' },
+    ])
+  })
+
+  conditionalTest('object contents input', async () => {
+    await client.models.generateContent({
+      model: 'gemini-2.0-flash-001',
+      contents: { text: 'Hello world' },
+      posthogDistinctId: 'test-id',
+    })
+
+    expect(mockPostHogClient.capture).toHaveBeenCalledTimes(1)
+    const [captureArgs] = (mockPostHogClient.capture as jest.Mock).mock.calls
+    const { properties } = captureArgs[0]
+
+    expect(properties['$ai_input']).toEqual([{ role: 'user', content: 'Hello world' }])
+  })
+
+  conditionalTest('capture immediate', async () => {
+    await client.models.generateContent({
+      model: 'gemini-2.0-flash-001',
+      contents: 'Hello',
+      posthogDistinctId: 'test-id',
+      posthogCaptureImmediate: true,
+    })
+
+    expect(mockPostHogClient.captureImmediate).toHaveBeenCalledTimes(1)
+    expect(mockPostHogClient.capture).toHaveBeenCalledTimes(0)
+  })
+
+  conditionalTest('vertex ai configuration', () => {
+    const vertexClient = new PostHogGemini({
+      vertexai: true,
+      project: 'test-project',
+      location: 'us-central1',
+      posthog: mockPostHogClient as any,
+    })
+
+    expect(vertexClient).toBeInstanceOf(PostHogGemini)
+    expect(vertexClient.models).toBeDefined()
+  })
+})