@posthog/ai 4.2.0 → 4.3.0

package/src/openai/index.ts

@@ -76,6 +76,7 @@ export class WrappedCompletions extends OpenAIOrignal.Chat.Completions {
       // eslint-disable-next-line @typescript-eslint/no-unused-vars
       posthogPrivacyMode = false,
       posthogGroups,
+      posthogCaptureImmediate,
       ...openAIParams
     } = body
 
@@ -115,7 +116,7 @@ export class WrappedCompletions extends OpenAIOrignal.Chat.Completions {
               }
 
               const latency = (Date.now() - startTime) / 1000
-              sendEventToPosthog({
+              await sendEventToPosthog({
                 client: this.phClient,
                 distinctId: posthogDistinctId ?? traceId,
                 traceId,
@@ -128,9 +129,10 @@ export class WrappedCompletions extends OpenAIOrignal.Chat.Completions {
                 params: body,
                 httpStatus: 200,
                 usage,
+                captureImmediate: posthogCaptureImmediate,
               })
             } catch (error: any) {
-              sendEventToPosthog({
+              await sendEventToPosthog({
                 client: this.phClient,
                 distinctId: posthogDistinctId ?? traceId,
                 traceId,
@@ -145,6 +147,7 @@ export class WrappedCompletions extends OpenAIOrignal.Chat.Completions {
                 usage: { inputTokens: 0, outputTokens: 0 },
                 isError: true,
                 error: JSON.stringify(error),
+                captureImmediate: posthogCaptureImmediate,
               })
             }
           })()
@@ -156,10 +159,10 @@ export class WrappedCompletions extends OpenAIOrignal.Chat.Completions {
       }) as APIPromise<Stream<ChatCompletionChunk>>
     } else {
       const wrappedPromise = parentPromise.then(
-        (result) => {
+        async (result) => {
           if ('choices' in result) {
             const latency = (Date.now() - startTime) / 1000
-            sendEventToPosthog({
+            await sendEventToPosthog({
               client: this.phClient,
               distinctId: posthogDistinctId ?? traceId,
               traceId,
@@ -177,12 +180,13 @@ export class WrappedCompletions extends OpenAIOrignal.Chat.Completions {
                 reasoningTokens: result.usage?.completion_tokens_details?.reasoning_tokens ?? 0,
                 cacheReadInputTokens: result.usage?.prompt_tokens_details?.cached_tokens ?? 0,
               },
+              captureImmediate: posthogCaptureImmediate,
             })
           }
           return result
         },
-        (error: any) => {
-          sendEventToPosthog({
+        async (error: any) => {
+          await sendEventToPosthog({
             client: this.phClient,
             distinctId: posthogDistinctId ?? traceId,
             traceId,
@@ -200,6 +204,7 @@ export class WrappedCompletions extends OpenAIOrignal.Chat.Completions {
             },
             isError: true,
             error: JSON.stringify(error),
+            captureImmediate: posthogCaptureImmediate,
           })
           throw error
         }

package/src/utils.ts

@@ -19,7 +19,7 @@ export interface MonitoringParams {
   posthogModelOverride?: string
   posthogProviderOverride?: string
   posthogCostOverride?: CostOverride
-  fullDebug?: boolean
+  posthogCaptureImmediate?: boolean
 }
 
 export interface CostOverride {
@@ -149,7 +149,7 @@ export type SendEventToPosthogParams = {
   isError?: boolean
   error?: string
   tools?: any
-  fullDebug?: boolean
+  captureImmediate?: boolean
 }
 
 function sanitizeValues(obj: any): any {
@@ -167,7 +167,7 @@ function sanitizeValues(obj: any): any {
   return jsonSafe
 }
 
-export const sendEventToPosthog = ({
+export const sendEventToPosthog = async ({
   client,
   distinctId,
   traceId,
@@ -183,68 +183,69 @@ export const sendEventToPosthog = ({
   isError = false,
   error,
   tools,
-  fullDebug = false,
-}: SendEventToPosthogParams): void => {
-  if (client.capture) {
-    // sanitize input and output for UTF-8 validity
-    const safeInput = sanitizeValues(input)
-    const safeOutput = sanitizeValues(output)
-    const safeError = sanitizeValues(error)
-
-    let errorData = {}
-    if (isError) {
-      errorData = {
-        $ai_is_error: true,
-        $ai_error: safeError,
-      }
+  captureImmediate = false,
+}: SendEventToPosthogParams): Promise<void> => {
+  if (!client.capture) return Promise.resolve()
+  // sanitize input and output for UTF-8 validity
+  const safeInput = sanitizeValues(input)
+  const safeOutput = sanitizeValues(output)
+  const safeError = sanitizeValues(error)
+
+  let errorData = {}
+  if (isError) {
+    errorData = {
+      $ai_is_error: true,
+      $ai_error: safeError,
     }
-    let costOverrideData = {}
-    if (params.posthogCostOverride) {
-      const inputCostUSD = (params.posthogCostOverride.inputCost ?? 0) * (usage.inputTokens ?? 0)
-      const outputCostUSD = (params.posthogCostOverride.outputCost ?? 0) * (usage.outputTokens ?? 0)
-      costOverrideData = {
-        $ai_input_cost_usd: inputCostUSD,
-        $ai_output_cost_usd: outputCostUSD,
-        $ai_total_cost_usd: inputCostUSD + outputCostUSD,
-      }
+  }
+  let costOverrideData = {}
+  if (params.posthogCostOverride) {
+    const inputCostUSD = (params.posthogCostOverride.inputCost ?? 0) * (usage.inputTokens ?? 0)
+    const outputCostUSD = (params.posthogCostOverride.outputCost ?? 0) * (usage.outputTokens ?? 0)
+    costOverrideData = {
+      $ai_input_cost_usd: inputCostUSD,
+      $ai_output_cost_usd: outputCostUSD,
+      $ai_total_cost_usd: inputCostUSD + outputCostUSD,
     }
+  }
 
-    const additionalTokenValues = {
-      ...(usage.reasoningTokens ? { $ai_reasoning_tokens: usage.reasoningTokens } : {}),
-      ...(usage.cacheReadInputTokens ? { $ai_cache_read_input_tokens: usage.cacheReadInputTokens } : {}),
-      ...(usage.cacheCreationInputTokens ? { $ai_cache_creation_input_tokens: usage.cacheCreationInputTokens } : {}),
-    }
+  const additionalTokenValues = {
+    ...(usage.reasoningTokens ? { $ai_reasoning_tokens: usage.reasoningTokens } : {}),
+    ...(usage.cacheReadInputTokens ? { $ai_cache_read_input_tokens: usage.cacheReadInputTokens } : {}),
+    ...(usage.cacheCreationInputTokens ? { $ai_cache_creation_input_tokens: usage.cacheCreationInputTokens } : {}),
+  }
 
-    const properties = {
-      $ai_provider: params.posthogProviderOverride ?? provider,
-      $ai_model: params.posthogModelOverride ?? model,
-      $ai_model_parameters: getModelParams(params),
-      $ai_input: withPrivacyMode(client, params.posthogPrivacyMode ?? false, safeInput),
-      $ai_output_choices: withPrivacyMode(client, params.posthogPrivacyMode ?? false, safeOutput),
-      $ai_http_status: httpStatus,
-      $ai_input_tokens: usage.inputTokens ?? 0,
-      $ai_output_tokens: usage.outputTokens ?? 0,
-      ...additionalTokenValues,
-      $ai_latency: latency,
-      $ai_trace_id: traceId,
-      $ai_base_url: baseURL,
-      ...params.posthogProperties,
-      ...(distinctId ? {} : { $process_person_profile: false }),
-      ...(tools ? { $ai_tools: tools } : {}),
-      ...errorData,
-      ...costOverrideData,
-    }
+  const properties = {
+    $ai_provider: params.posthogProviderOverride ?? provider,
+    $ai_model: params.posthogModelOverride ?? model,
+    $ai_model_parameters: getModelParams(params),
+    $ai_input: withPrivacyMode(client, params.posthogPrivacyMode ?? false, safeInput),
+    $ai_output_choices: withPrivacyMode(client, params.posthogPrivacyMode ?? false, safeOutput),
+    $ai_http_status: httpStatus,
+    $ai_input_tokens: usage.inputTokens ?? 0,
+    $ai_output_tokens: usage.outputTokens ?? 0,
+    ...additionalTokenValues,
+    $ai_latency: latency,
+    $ai_trace_id: traceId,
+    $ai_base_url: baseURL,
+    ...params.posthogProperties,
+    ...(distinctId ? {} : { $process_person_profile: false }),
+    ...(tools ? { $ai_tools: tools } : {}),
+    ...errorData,
+    ...costOverrideData,
+  }
 
-    if (fullDebug) {
-      // @ts-ignore
-      console.log('Sending event to PostHog', properties)
-    }
+  const event = {
+    distinctId: distinctId ?? traceId,
+    event: '$ai_generation',
+    properties,
+    groups: params.posthogGroups,
+  }
 
-    client.capture({
-      distinctId: distinctId ?? traceId,
-      event: '$ai_generation',
-      properties,
-      groups: params.posthogGroups,
-    })
+  if (captureImmediate) {
+    // await capture promise to send single event in serverless environments
+    await client.captureImmediate(event)
+  } else {
+    client.capture(event)
   }
 }

package/src/vercel/middleware.ts

@@ -14,7 +14,7 @@ interface ClientOptions {
   posthogModelOverride?: string
   posthogProviderOverride?: string
   posthogCostOverride?: CostOverride
-  fullDebug?: boolean
+  posthogCaptureImmediate?: boolean
 }
 
 interface CreateInstrumentationMiddlewareOptions {
@@ -26,7 +26,7 @@ interface CreateInstrumentationMiddlewareOptions {
   posthogModelOverride?: string
   posthogProviderOverride?: string
   posthogCostOverride?: CostOverride
-  fullDebug?: boolean
+  posthogCaptureImmediate?: boolean
 }
 
 interface PostHogInput {
@@ -127,15 +127,23 @@ const mapVercelPrompt = (prompt: LanguageModelV1Prompt): PostHogInput[] => {
   try {
     // Trim the inputs array until its JSON size fits within MAX_OUTPUT_SIZE
     let serialized = JSON.stringify(inputs)
-    while (Buffer.byteLength(serialized, 'utf8') > MAX_OUTPUT_SIZE && inputs.length > 0) {
-      // Remove oldest message
+    let removedCount = 0
+    // We need to keep track of the initial size of the inputs array because we're going to be mutating it
+    let initialSize = inputs.length
+    for (let i = 0; i < initialSize && Buffer.byteLength(serialized, 'utf8') > MAX_OUTPUT_SIZE; i++) {
       inputs.shift()
-      // add blank message to beginning of array
-      inputs.unshift({ role: 'assistant', content: '[removed message due to size limit]' })
+      removedCount++
       serialized = JSON.stringify(inputs)
     }
+    if (removedCount > 0) {
+      // Add one placeholder to indicate how many were removed
+      inputs.unshift({
+        role: 'posthog',
+        content: `[${removedCount} message${removedCount === 1 ? '' : 's'} removed due to size limit]`,
+      })
+    }
   } catch (error) {
-    console.error('Error stringifying inputs')
+    console.error('Error stringifying inputs', error)
     return [{ role: 'posthog', content: 'An error occurred while processing your request. Please try again.' }]
   }
   return inputs
@@ -218,7 +226,7 @@ export const createInstrumentationMiddleware = (
               }
             : {}),
         }
-        sendEventToPosthog({
+        await sendEventToPosthog({
           client: phClient,
           distinctId: options.posthogDistinctId,
           traceId: options.posthogTraceId,
@@ -235,13 +243,13 @@ export const createInstrumentationMiddleware = (
             outputTokens: result.usage.completionTokens,
             ...additionalTokenValues,
           },
-          fullDebug: options.fullDebug,
+          captureImmediate: options.posthogCaptureImmediate,
         })
 
         return result
       } catch (error: any) {
         const modelId = model.modelId
-        sendEventToPosthog({
+        await sendEventToPosthog({
           client: phClient,
           distinctId: options.posthogDistinctId,
           traceId: options.posthogTraceId,
@@ -259,7 +267,7 @@ export const createInstrumentationMiddleware = (
           },
           isError: true,
           error: truncate(JSON.stringify(error)),
-          fullDebug: options.fullDebug,
+          captureImmediate: options.posthogCaptureImmediate,
         })
         throw error
       }
@@ -311,9 +319,9 @@ export const createInstrumentationMiddleware = (
             controller.enqueue(chunk)
           },
 
-          flush() {
+          flush: async () => {
             const latency = (Date.now() - startTime) / 1000
-            sendEventToPosthog({
+            await sendEventToPosthog({
               client: phClient,
               distinctId: options.posthogDistinctId,
               traceId: options.posthogTraceId,
@@ -326,7 +334,7 @@ export const createInstrumentationMiddleware = (
               params: mergedParams as any,
               httpStatus: 200,
               usage,
-              fullDebug: options.fullDebug,
+              captureImmediate: options.posthogCaptureImmediate,
             })
           },
         })
@@ -336,7 +344,7 @@ export const createInstrumentationMiddleware = (
           ...rest,
         }
       } catch (error: any) {
-        sendEventToPosthog({
+        await sendEventToPosthog({
           client: phClient,
           distinctId: options.posthogDistinctId,
           traceId: options.posthogTraceId,
@@ -354,7 +362,7 @@ export const createInstrumentationMiddleware = (
           },
           isError: true,
           error: truncate(JSON.stringify(error)),
-          fullDebug: options.fullDebug,
+          captureImmediate: options.posthogCaptureImmediate,
         })
         throw error
       }

package/tests/openai.test.ts

@@ -6,6 +6,7 @@ jest.mock('posthog-node', () => {
     PostHog: jest.fn().mockImplementation(() => {
       return {
         capture: jest.fn(),
+        captureImmediate: jest.fn(),
         privacyMode: false,
       }
     }),
@@ -261,4 +262,18 @@ describe('PostHogOpenAI - Jest test suite', () => {
     expect(properties['$ai_reasoning_tokens']).toBe(15)
     expect(properties['$ai_cache_read_input_tokens']).toBe(5)
   })
+
+  // New test: ensure captureImmediate is used when flag is set
+  conditionalTest('captureImmediate flag', async () => {
+    await client.chat.completions.create({
+      model: 'gpt-4',
+      messages: [{ role: 'user', content: 'Hello' }],
+      posthogDistinctId: 'test-id',
+      posthogCaptureImmediate: true,
+    })
+
+    // captureImmediate should be called once, and capture should not be called
+    expect(mockPostHogClient.captureImmediate).toHaveBeenCalledTimes(1)
+    expect(mockPostHogClient.capture).toHaveBeenCalledTimes(0)
+  })
 })

package/tsconfig.json

@@ -2,7 +2,7 @@
   "extends": "../tsconfig.json",
   "compilerOptions": {
     "incremental": false,
-    "types": ["node"],
+    "types": ["node", "jest"],
     "typeRoots": ["./node_modules/@types", "../node_modules/@types"],
     "moduleResolution": "node",
     "skipLibCheck": true,