@posthog/ai 1.0.1

package/README.md

@@ -0,0 +1,62 @@
+# PostHog Node AI
+
+Initial Typescript SDK for LLM Observability
+
+## Installation
+
+```bash
+yarn add @posthog/posthog-ai
+```
+
+## Usage
+
+### Before
+
+```typescript
+import { OpenAI } from 'openai'
+
+const client = new OpenAI({
+  apiKey: process.env.OPENAI_API_KEY || '',
+})
+
+await client.chat.completions.create({
+  model: 'gpt-4',
+  messages: [{ role: 'user', content: 'Hello, world!' }],
+})
+```
+
+### After
+
+```typescript
+import { OpenAI } from 'posthog-node-ai'
+import { PostHog } from 'posthog-node'
+
+const phClient = new PostHog(
+  process.env.POSTHOG_API_KEY, {
+    host: process.env.POSTHOG_HOST || 'https://us.posthog.com',
+  }
+})
+
+const client = new OpenAI({
+  apiKey: process.env.OPENAI_API_KEY || '',
+  posthog: phClient,
+})
+
+await client.chat.completions.create({
+  model: 'gpt-4',
+  messages: [{ role: 'user', content: 'Hello, world!' }],
+  posthog_distinct_id: 'test-user-id',
+  posthog_properties: {
+    test_property: 'test_value',
+  }
+})
+
+// YOU HAVE TO HAVE THIS OR THE CLIENT MAY NOT SEND EVENTS
+await phClient.shutdown()
+```
+
+Please see the main [PostHog docs](https://www.posthog.com/docs).
+
+## Questions?
+
+### [Check out our community page.](https://posthog.com/posts)

package/index.ts

@@ -0,0 +1 @@
+export * from './src'

package/package.json

@@ -0,0 +1,48 @@
+{
+  "name": "@posthog/ai",
+  "version": "1.0.1",
+  "description": "PostHog Node.js AI integrations",
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/PostHog/posthog-js-lite.git",
+    "directory": "posthog-ai"
+  },
+  "main": "./lib/index.js",
+  "module": "./lib/index.mjs",
+  "types": "./lib/index.d.ts",
+  "license": "MIT",
+  "devDependencies": {
+    "@types/jest": "^28.1.5",
+    "@types/node": "^18.0.0",
+    "jest": "^29.0.0",
+    "node-fetch": "^3.3.2",
+    "ts-jest": "^29.0.0",
+    "typescript": "^4.7.4"
+  },
+  "keywords": [
+    "posthog",
+    "ai",
+    "openai",
+    "anthropic",
+    "llm"
+  ],
+  "dependencies": {
+    "ai": "^4.1.0",
+    "openai": "^4.79.1",
+    "uuid": "^11.0.5",
+    "zod": "^3.24.1"
+  },
+  "scripts": {
+    "test": "jest",
+    "prepublishOnly": "cd .. && yarn build"
+  },
+  "directories": {
+    "lib": "lib",
+    "test": "tests"
+  },
+  "author": "PostHog",
+  "bugs": {
+    "url": "https://github.com/PostHog/posthog-js-lite/issues"
+  },
+  "homepage": "https://github.com/PostHog/posthog-js-lite#readme"
+}

package/src/azure-openai/index.ts

@@ -0,0 +1,204 @@
+import OpenAIOrignal, { AzureOpenAI } from 'openai'
+import { PostHog } from 'posthog-node'
+import { v4 as uuidv4 } from 'uuid'
+import { PassThrough } from 'stream'
+import { mergeSystemPrompt, type MonitoringParams, sendEventToPosthog } from '../utils'
+
+type ChatCompletion = OpenAIOrignal.ChatCompletion
+type ChatCompletionChunk = OpenAIOrignal.ChatCompletionChunk
+type ChatCompletionCreateParamsBase = OpenAIOrignal.Chat.Completions.ChatCompletionCreateParams
+type ChatCompletionCreateParamsNonStreaming = OpenAIOrignal.Chat.Completions.ChatCompletionCreateParamsNonStreaming
+type ChatCompletionCreateParamsStreaming = OpenAIOrignal.Chat.Completions.ChatCompletionCreateParamsStreaming
+import type { APIPromise, RequestOptions } from 'openai/core'
+import type { Stream } from 'openai/streaming'
+
+interface MonitoringOpenAIConfig {
+  apiKey: string
+  posthog: PostHog
+  baseURL?: string
+}
+
+export class PostHogAzureOpenAI extends AzureOpenAI {
+  private readonly phClient: PostHog
+
+  constructor(config: MonitoringOpenAIConfig) {
+    const { posthog, ...openAIConfig } = config
+    super(openAIConfig)
+    this.phClient = posthog
+    this.chat = new WrappedChat(this, this.phClient)
+  }
+
+  public chat: WrappedChat
+}
+
+export class WrappedChat extends AzureOpenAI.Chat {
+  constructor(parentClient: PostHogAzureOpenAI, phClient: PostHog) {
+    super(parentClient)
+    this.completions = new WrappedCompletions(parentClient, phClient)
+  }
+
+  public completions: WrappedCompletions
+}
+
+export class WrappedCompletions extends AzureOpenAI.Chat.Completions {
+  private readonly phClient: PostHog
+
+  constructor(client: AzureOpenAI, phClient: PostHog) {
+    super(client)
+    this.phClient = phClient
+  }
+
+  // --- Overload #1: Non-streaming
+  public create(
+    body: ChatCompletionCreateParamsNonStreaming & MonitoringParams,
+    options?: RequestOptions
+  ): APIPromise<ChatCompletion>
+
+  // --- Overload #2: Streaming
+  public create(
+    body: ChatCompletionCreateParamsStreaming & MonitoringParams,
+    options?: RequestOptions
+  ): APIPromise<Stream<ChatCompletionChunk>>
+
+  // --- Overload #3: Generic base
+  public create(
+    body: ChatCompletionCreateParamsBase & MonitoringParams,
+    options?: RequestOptions
+  ): APIPromise<ChatCompletion | Stream<ChatCompletionChunk>>
+
+  // --- Implementation Signature
+  public create(
+    body: ChatCompletionCreateParamsBase & MonitoringParams,
+    options?: RequestOptions
+  ): APIPromise<ChatCompletion | Stream<ChatCompletionChunk>> {
+    const {
+      posthog_distinct_id,
+      posthog_trace_id,
+      posthog_properties,
+      posthog_privacy_mode = false,
+      posthog_groups,
+      ...openAIParams
+    } = body
+
+    const traceId = posthog_trace_id ?? uuidv4()
+    const startTime = Date.now()
+
+    const parentPromise = super.create(openAIParams, options)
+
+    if (openAIParams.stream) {
+      return parentPromise.then((value) => {
+        const passThroughStream = new PassThrough({ objectMode: true })
+        let accumulatedContent = ''
+        let usage: { input_tokens: number; output_tokens: number } = {
+          input_tokens: 0,
+          output_tokens: 0,
+        }
+        if ('tee' in value) {
+          const openAIStream = value
+          ;(async () => {
+            try {
+              for await (const chunk of openAIStream) {
+                const delta = chunk?.choices?.[0]?.delta?.content ?? ''
+                accumulatedContent += delta
+                if (chunk.usage) {
+                  usage = {
+                    input_tokens: chunk.usage.prompt_tokens ?? 0,
+                    output_tokens: chunk.usage.completion_tokens ?? 0,
+                  }
+                }
+                passThroughStream.write(chunk)
+              }
+              const latency = (Date.now() - startTime) / 1000
+              sendEventToPosthog({
+                client: this.phClient,
+                distinctId: posthog_distinct_id ?? traceId,
+                traceId,
+                model: openAIParams.model,
+                provider: 'azure',
+                input: posthog_privacy_mode ? '' : mergeSystemPrompt(openAIParams, 'openai'),
+                output: [{ content: accumulatedContent, role: 'assistant' }],
+                latency,
+                baseURL: (this as any).baseURL ?? '',
+                params: body,
+                httpStatus: 200,
+                usage,
+              })
+              passThroughStream.end()
+            } catch (error) {
+              // error handling
+              sendEventToPosthog({
+                client: this.phClient,
+                distinctId: posthog_distinct_id ?? traceId,
+                traceId,
+                model: openAIParams.model,
+                provider: 'azure',
+                input: posthog_privacy_mode ? '' : mergeSystemPrompt(openAIParams, 'openai'),
+                output: [],
+                latency: 0,
+                baseURL: (this as any).baseURL ?? '',
+                params: body,
+                httpStatus: 500,
+                usage: {
+                  input_tokens: 0,
+                  output_tokens: 0,
+                },
+              })
+              passThroughStream.emit('error', error)
+            }
+          })()
+        }
+        return passThroughStream as unknown as Stream<ChatCompletionChunk>
+      }) as APIPromise<Stream<ChatCompletionChunk>>
+    } else {
+      const wrappedPromise = parentPromise.then(
+        (result) => {
+          if ('choices' in result) {
+            const latency = (Date.now() - startTime) / 1000
+            sendEventToPosthog({
+              client: this.phClient,
+              distinctId: posthog_distinct_id ?? traceId,
+              traceId,
+              model: openAIParams.model,
+              provider: 'azure',
+              input: posthog_privacy_mode ? '' : mergeSystemPrompt(openAIParams, 'openai'),
+              output: [{ content: result.choices[0].message.content, role: 'assistant' }],
+              latency,
+              baseURL: (this as any).baseURL ?? '',
+              params: body,
+              httpStatus: 200,
+              usage: {
+                input_tokens: result.usage?.prompt_tokens ?? 0,
+                output_tokens: result.usage?.completion_tokens ?? 0,
+              },
+            })
+          }
+          return result
+        },
+        (error) => {
+          sendEventToPosthog({
+            client: this.phClient,
+            distinctId: posthog_distinct_id ?? traceId,
+            traceId,
+            model: openAIParams.model,
+            provider: 'azure',
+            input: posthog_privacy_mode ? '' : mergeSystemPrompt(openAIParams, 'openai'),
+            output: [],
+            latency: 0,
+            baseURL: (this as any).baseURL ?? '',
+            params: body,
+            httpStatus: 500,
+            usage: {
+              input_tokens: 0,
+              output_tokens: 0,
+            },
+          })
+          throw error
+        }
+      ) as APIPromise<ChatCompletion>
+
+      return wrappedPromise
+    }
+  }
+}
+
+export default PostHogAzureOpenAI

package/src/index.ts

@@ -0,0 +1,7 @@
+import PostHogOpenAI from './openai'
+import PostHogAzureOpenAI from './azure-openai'
+import { wrapVercelLanguageModel } from './vercel/middleware'
+
+export { PostHogOpenAI as OpenAI }
+export { PostHogAzureOpenAI as AzureOpenAI }
+export { wrapVercelLanguageModel as posthogWrappedLanguageModel }

package/src/openai/index.ts

@@ -0,0 +1,204 @@
+import OpenAIOrignal from 'openai'
+import { PostHog } from 'posthog-node'
+import { v4 as uuidv4 } from 'uuid'
+import { PassThrough } from 'stream'
+import { mergeSystemPrompt, type MonitoringParams, sendEventToPosthog } from '../utils'
+
+type ChatCompletion = OpenAIOrignal.ChatCompletion
+type ChatCompletionChunk = OpenAIOrignal.ChatCompletionChunk
+type ChatCompletionCreateParamsBase = OpenAIOrignal.Chat.Completions.ChatCompletionCreateParams
+type ChatCompletionCreateParamsNonStreaming = OpenAIOrignal.Chat.Completions.ChatCompletionCreateParamsNonStreaming
+type ChatCompletionCreateParamsStreaming = OpenAIOrignal.Chat.Completions.ChatCompletionCreateParamsStreaming
+import type { APIPromise, RequestOptions } from 'openai/core'
+import type { Stream } from 'openai/streaming'
+
+interface MonitoringOpenAIConfig {
+  apiKey: string
+  posthog: PostHog
+  baseURL?: string
+}
+
+export class PostHogOpenAI extends OpenAIOrignal {
+  private readonly phClient: PostHog
+
+  constructor(config: MonitoringOpenAIConfig) {
+    const { posthog, ...openAIConfig } = config
+    super(openAIConfig)
+    this.phClient = posthog
+    this.chat = new WrappedChat(this, this.phClient)
+  }
+
+  public chat: WrappedChat
+}
+
+export class WrappedChat extends OpenAIOrignal.Chat {
+  constructor(parentClient: PostHogOpenAI, phClient: PostHog) {
+    super(parentClient)
+    this.completions = new WrappedCompletions(parentClient, phClient)
+  }
+
+  public completions: WrappedCompletions
+}
+
+export class WrappedCompletions extends OpenAIOrignal.Chat.Completions {
+  private readonly phClient: PostHog
+
+  constructor(client: OpenAIOrignal, phClient: PostHog) {
+    super(client)
+    this.phClient = phClient
+  }
+
+  // --- Overload #1: Non-streaming
+  public create(
+    body: ChatCompletionCreateParamsNonStreaming & MonitoringParams,
+    options?: RequestOptions
+  ): APIPromise<ChatCompletion>
+
+  // --- Overload #2: Streaming
+  public create(
+    body: ChatCompletionCreateParamsStreaming & MonitoringParams,
+    options?: RequestOptions
+  ): APIPromise<Stream<ChatCompletionChunk>>
+
+  // --- Overload #3: Generic base
+  public create(
+    body: ChatCompletionCreateParamsBase & MonitoringParams,
+    options?: RequestOptions
+  ): APIPromise<ChatCompletion | Stream<ChatCompletionChunk>>
+
+  // --- Implementation Signature
+  public create(
+    body: ChatCompletionCreateParamsBase & MonitoringParams,
+    options?: RequestOptions
+  ): APIPromise<ChatCompletion | Stream<ChatCompletionChunk>> {
+    const {
+      posthog_distinct_id,
+      posthog_trace_id,
+      posthog_properties,
+      posthog_privacy_mode = false,
+      posthog_groups,
+      ...openAIParams
+    } = body
+
+    const traceId = posthog_trace_id ?? uuidv4()
+    const startTime = Date.now()
+
+    const parentPromise = super.create(openAIParams, options)
+
+    if (openAIParams.stream) {
+      return parentPromise.then((value) => {
+        const passThroughStream = new PassThrough({ objectMode: true })
+        let accumulatedContent = ''
+        let usage: { input_tokens: number; output_tokens: number } = {
+          input_tokens: 0,
+          output_tokens: 0,
+        }
+        if ('tee' in value) {
+          const openAIStream = value
+          ;(async () => {
+            try {
+              for await (const chunk of openAIStream) {
+                const delta = chunk?.choices?.[0]?.delta?.content ?? ''
+                accumulatedContent += delta
+                if (chunk.usage) {
+                  usage = {
+                    input_tokens: chunk.usage.prompt_tokens ?? 0,
+                    output_tokens: chunk.usage.completion_tokens ?? 0,
+                  }
+                }
+                passThroughStream.write(chunk)
+              }
+              const latency = (Date.now() - startTime) / 1000
+              sendEventToPosthog({
+                client: this.phClient,
+                distinctId: posthog_distinct_id ?? traceId,
+                traceId,
+                model: openAIParams.model,
+                provider: 'openai',
+                input: posthog_privacy_mode ? '' : mergeSystemPrompt(openAIParams, 'openai'),
+                output: [{ content: accumulatedContent, role: 'assistant' }],
+                latency,
+                baseURL: (this as any).baseURL ?? '',
+                params: body,
+                httpStatus: 200,
+                usage,
+              })
+              passThroughStream.end()
+            } catch (error) {
+              // error handling
+              sendEventToPosthog({
+                client: this.phClient,
+                distinctId: posthog_distinct_id ?? traceId,
+                traceId,
+                model: openAIParams.model,
+                provider: 'openai',
+                input: posthog_privacy_mode ? '' : mergeSystemPrompt(openAIParams, 'openai'),
+                output: [],
+                latency: 0,
+                baseURL: (this as any).baseURL ?? '',
+                params: body,
+                httpStatus: 500,
+                usage: {
+                  input_tokens: 0,
+                  output_tokens: 0,
+                },
+              })
+              passThroughStream.emit('error', error)
+            }
+          })()
+        }
+        return passThroughStream as unknown as Stream<ChatCompletionChunk>
+      }) as APIPromise<Stream<ChatCompletionChunk>>
+    } else {
+      const wrappedPromise = parentPromise.then(
+        (result) => {
+          if ('choices' in result) {
+            const latency = (Date.now() - startTime) / 1000
+            sendEventToPosthog({
+              client: this.phClient,
+              distinctId: posthog_distinct_id ?? traceId,
+              traceId,
+              model: openAIParams.model,
+              provider: 'openai',
+              input: posthog_privacy_mode ? '' : mergeSystemPrompt(openAIParams, 'openai'),
+              output: [{ content: result.choices[0].message.content, role: 'assistant' }],
+              latency,
+              baseURL: (this as any).baseURL ?? '',
+              params: body,
+              httpStatus: 200,
+              usage: {
+                input_tokens: result.usage?.prompt_tokens ?? 0,
+                output_tokens: result.usage?.completion_tokens ?? 0,
+              },
+            })
+          }
+          return result
+        },
+        (error) => {
+          sendEventToPosthog({
+            client: this.phClient,
+            distinctId: posthog_distinct_id ?? traceId,
+            traceId,
+            model: openAIParams.model,
+            provider: 'openai',
+            input: posthog_privacy_mode ? '' : mergeSystemPrompt(openAIParams, 'openai'),
+            output: [],
+            latency: 0,
+            baseURL: (this as any).baseURL ?? '',
+            params: body,
+            httpStatus: 500,
+            usage: {
+              input_tokens: 0,
+              output_tokens: 0,
+            },
+          })
+          throw error
+        }
+      ) as APIPromise<ChatCompletion>
+
+      return wrappedPromise
+    }
+  }
+}
+
+export default PostHogOpenAI

package/src/utils.ts

@@ -0,0 +1,166 @@
+import { PostHog } from 'posthog-node'
+import OpenAIOrignal from 'openai'
+
+type ChatCompletionCreateParamsBase = OpenAIOrignal.Chat.Completions.ChatCompletionCreateParams
+
+export interface MonitoringParams {
+  posthog_distinct_id?: string
+  posthog_trace_id?: string
+  posthog_properties?: Record<string, any>
+  posthog_privacy_mode?: boolean
+  posthog_groups?: Record<string, any>
+}
+
+export const getModelParams = (params: ChatCompletionCreateParamsBase & MonitoringParams): Record<string, any> => {
+  const modelParams: Record<string, any> = {}
+  const paramKeys = [
+    'temperature',
+    'max_tokens',
+    'max_completion_tokens',
+    'top_p',
+    'frequency_penalty',
+    'presence_penalty',
+    'n',
+    'stop',
+    'stream',
+    'streaming',
+  ] as const
+
+  for (const key of paramKeys) {
+    if (key in params && (params as any)[key] !== undefined) {
+      modelParams[key] = (params as any)[key]
+    }
+  }
+  return modelParams
+}
+
+export const getUsage = (response: any, provider: string): { input_tokens: number; output_tokens: number } => {
+  if (!response?.usage) {
+    return { input_tokens: 0, output_tokens: 0 }
+  }
+
+  if (provider === 'anthropic') {
+    return {
+      input_tokens: response.usage.input_tokens ?? 0,
+      output_tokens: response.usage.output_tokens ?? 0,
+    }
+  } else if (provider === 'openai') {
+    return {
+      input_tokens: response.usage.prompt_tokens ?? 0,
+      output_tokens: response.usage.completion_tokens ?? 0,
+    }
+  }
+
+  return { input_tokens: 0, output_tokens: 0 }
+}
+
+/**
+ * Helper to format responses (non-streaming) for consumption, mirroring Python's openai vs. anthropic approach.
+ */
+export const formatResponse = (response: any, provider: string): Array<{ role: string; content: string }> => {
+  if (!response) {
+    return []
+  }
+  if (provider === 'anthropic') {
+    return formatResponseAnthropic(response)
+  } else if (provider === 'openai') {
+    return formatResponseOpenAI(response)
+  }
+  return []
+}
+
+export const formatResponseAnthropic = (response: any): Array<{ role: string; content: string }> => {
+  // Example approach if "response.content" holds array of text segments, etc.
+  const output: Array<{ role: string; content: string }> = []
+  for (const choice of response.content ?? []) {
+    if (choice?.text) {
+      output.push({
+        role: 'assistant',
+        content: choice.text,
+      })
+    }
+  }
+  return output
+}
+
+export const formatResponseOpenAI = (response: any): Array<{ role: string; content: string }> => {
+  const output: Array<{ role: string; content: string }> = []
+  for (const choice of response.choices ?? []) {
+    if (choice.message?.content) {
+      output.push({
+        role: choice.message.role,
+        content: choice.message.content,
+      })
+    }
+  }
+  return output
+}
+
+export const mergeSystemPrompt = (params: ChatCompletionCreateParamsBase & MonitoringParams, provider: string): any => {
+  if (provider !== 'anthropic') {
+    return params.messages
+  }
+  const messages = params.messages || []
+  if (!(params as any).system) {
+    return messages
+  }
+  const systemMessage = (params as any).system
+  return [{ role: 'system', content: systemMessage }, ...messages]
+}
+
+export const withPrivacyMode = (client: PostHog, privacyMode: boolean, input: any): any => {
+  return (client as any).privacy_mode || privacyMode ? null : input
+}
+
+export type SendEventToPosthogParams = {
+  client: PostHog
+  distinctId?: string
+  traceId: string
+  model: string
+  provider: string
+  input: any
+  output: any
+  latency: number
+  baseURL: string
+  httpStatus: number
+  usage: { input_tokens?: number; output_tokens?: number }
+  params: ChatCompletionCreateParamsBase & MonitoringParams
+}
+
+export const sendEventToPosthog = ({
+  client,
+  distinctId,
+  traceId,
+  model,
+  provider,
+  input,
+  output,
+  latency,
+  baseURL,
+  params,
+  httpStatus = 200,
+  usage = {},
+}: SendEventToPosthogParams): void => {
+  if (client.capture) {
+    client.capture({
+      distinctId: distinctId ?? traceId,
+      event: '$ai_generation',
+      properties: {
+        $ai_provider: provider,
+        $ai_model: model,
+        $ai_model_parameters: getModelParams(params),
+        $ai_input: withPrivacyMode(client, params.posthog_privacy_mode ?? false, input),
+        $ai_output_choices: withPrivacyMode(client, params.posthog_privacy_mode ?? false, output),
+        $ai_http_status: httpStatus,
+        $ai_input_tokens: usage.input_tokens ?? 0,
+        $ai_output_tokens: usage.output_tokens ?? 0,
+        $ai_latency: latency,
+        $ai_trace_id: traceId,
+        $ai_base_url: baseURL,
+        ...params.posthog_properties,
+        ...(distinctId ? {} : { $process_person_profile: false }),
+      },
+      groups: params.posthog_groups,
+    })
+  }
+}

package/src/vercel/middleware.ts

@@ -0,0 +1,157 @@
+import { experimental_wrapLanguageModel as wrapLanguageModel } from 'ai';
+import type {
+    LanguageModelV1,
+    Experimental_LanguageModelV1Middleware as LanguageModelV1Middleware,
+    LanguageModelV1StreamPart,
+} from 'ai';
+import { v4 as uuidv4 } from 'uuid';
+import type { PostHog } from 'posthog-node';
+import { sendEventToPosthog } from '../utils';
+
+interface CreateInstrumentationMiddlewareOptions {
+    posthog_distinct_id: string;
+    posthog_trace_id: string;
+    posthog_properties: Record<string, any>;
+    posthog_privacy_mode: boolean;
+    posthog_groups: string[];
+}
+
+export const createInstrumentationMiddleware = (phClient: PostHog, model: LanguageModelV1, options: CreateInstrumentationMiddlewareOptions) => {
+    const middleware: LanguageModelV1Middleware = {
+        wrapGenerate: async ({ doGenerate, params }) => {
+            const startTime = Date.now();
+
+            try {
+                const result = await doGenerate();
+                const latency = (Date.now() - startTime) / 1000;
+
+                sendEventToPosthog({
+                    client: phClient,
+                    distinctId: options.posthog_distinct_id,
+                    traceId: options.posthog_trace_id,
+                    model: model.modelId,
+                    provider: 'vercel',
+                    input: options.posthog_privacy_mode ? '' : params.prompt,
+                    output: [{ content: result.text, role: 'assistant' }],
+                    latency,
+                    baseURL: "",
+                    params: { posthog_properties: options } as any,
+                    httpStatus: 200,
+                    usage: {
+                        input_tokens: 0,
+                        output_tokens: 0,
+                    },
+                });
+
+                return result;
+            } catch (error) {
+                sendEventToPosthog({
+                    client: phClient,
+                    distinctId: options.posthog_distinct_id,
+                    traceId: options.posthog_trace_id,
+                    model: model.modelId,
+                    provider: 'vercel',
+                    input: options.posthog_privacy_mode ? '' : params.prompt,
+                    output: [],
+                    latency: 0,
+                    baseURL: "",
+                    params: { posthog_properties: options } as any,
+                    httpStatus: 500,
+                    usage: {
+                        input_tokens: 0,
+                        output_tokens: 0,
+                    },
+                });
+                throw error;
+            }
+        },
+
+        wrapStream: async ({ doStream, params }) => {
+            const startTime = Date.now();
+            let generatedText = '';
+
+            try {
+                const { stream, ...rest } = await doStream();
+
+                const transformStream = new TransformStream<
+                    LanguageModelV1StreamPart,
+                    LanguageModelV1StreamPart
+                >({
+                    transform(chunk, controller) {
+                        if (chunk.type === 'text-delta') {
+                            generatedText += chunk.textDelta;
+                        }
+                        controller.enqueue(chunk);
+                    },
+
+                    flush() {
+                        const latency = (Date.now() - startTime) / 1000;
+                        sendEventToPosthog({
+                            client: phClient,
+                            distinctId: options.posthog_distinct_id,
+                            traceId: options.posthog_trace_id,
+                            model: model.modelId,
+                            provider: 'vercel',
+                            input: options.posthog_privacy_mode ? '' : params.prompt,
+                            output: [{ content: generatedText, role: 'assistant' }],
+                            latency,
+                            baseURL: "",
+                            params: { posthog_properties: options } as any,
+                            httpStatus: 200,
+                            usage: {
+                                input_tokens: 0,
+                                output_tokens: 0,
+                            },
+                        });
+                    },
+                });
+
+                return {
+                    stream: stream.pipeThrough(transformStream),
+                    ...rest,
+                };
+            } catch (error) {
+                sendEventToPosthog({
+                    client: phClient,
+                    distinctId: options.posthog_distinct_id,
+                    traceId: options.posthog_trace_id,
+                    model: model.modelId,
+                    provider: 'vercel',
+                    input: options.posthog_privacy_mode ? '' : params.prompt,
+                    output: [],
+                    latency: 0,
+                    baseURL: "",
+                    params: { posthog_properties: options } as any,
+                    httpStatus: 500,
+                    usage: {
+                        input_tokens: 0,
+                        output_tokens: 0,
+                    },
+                });
+                throw error;
+            }
+        },
+    };
+
+    return middleware;
+};
+
+export const wrapVercelLanguageModel = (
+    model: LanguageModelV1,
+    phClient: PostHog,
+    options: CreateInstrumentationMiddlewareOptions
+) => {
+    const traceId = options.posthog_trace_id ?? uuidv4();
+    const middleware = createInstrumentationMiddleware(phClient, model, {
+        ...options,
+        posthog_trace_id: traceId,
+        posthog_distinct_id: options.posthog_distinct_id ?? traceId,
+    });
+
+    const wrappedModel = wrapLanguageModel({
+        model,
+        middleware,
+    });
+
+    return wrappedModel;
+};

package/tests/openai.test.ts

@@ -0,0 +1,228 @@
+import { PostHog } from 'posthog-node'
+import PostHogOpenAI from '../src/openai'
+
+jest.mock('posthog-node', () => {
+  return {
+    PostHog: jest.fn().mockImplementation(() => {
+      return {
+        capture: jest.fn(),
+        privacy_mode: false,
+      }
+    }),
+  }
+})
+
+let mockOpenAiChatResponse: any = {}
+let mockOpenAiEmbeddingResponse: any = {}
+
+describe('PostHogOpenAI - Jest test suite', () => {
+  let mockPostHogClient: PostHog
+  let client: PostHogOpenAI
+
+  beforeAll(() => {
+    if (!process.env.OPENAI_API_KEY) {
+      console.warn('⚠️ Skipping OpenAI tests: No OPENAI_API_KEY environment variable set')
+    }
+  })
+
+  beforeEach(() => {
+    // Skip all tests if no API key is present
+    if (!process.env.OPENAI_API_KEY) {
+      return
+    }
+
+    jest.clearAllMocks()
+
+    // Reset the default mocks
+    mockPostHogClient = new (PostHog as any)()
+    client = new PostHogOpenAI({
+      apiKey: process.env.OPENAI_API_KEY || '',
+      posthog: mockPostHogClient as any,
+    })
+
+    // Some default chat completion mock
+    mockOpenAiChatResponse = {
+      id: 'test-response-id',
+      model: 'gpt-4',
+      object: 'chat.completion',
+      created: Date.now() / 1000,
+      choices: [
+        {
+          index: 0,
+          finish_reason: 'stop',
+          message: {
+            role: 'assistant',
+            content: 'Hello from OpenAI!',
+          },
+        },
+      ],
+      usage: {
+        prompt_tokens: 20,
+        completion_tokens: 10,
+        total_tokens: 30,
+      },
+    }
+
+    // Some default embedding mock
+    mockOpenAiEmbeddingResponse = {
+      data: [
+        {
+          object: 'embedding',
+          index: 0,
+          embedding: [0.1, 0.2, 0.3],
+        },
+      ],
+      model: 'text-embedding-3-small',
+      object: 'list',
+      usage: {
+        prompt_tokens: 10,
+        total_tokens: 10,
+      },
+    }
+  })
+
+  // Wrap each test with conditional skip
+  const conditionalTest = process.env.OPENAI_API_KEY ? test : test.skip
+
+  conditionalTest('basic completion', async () => {
+    // We ensure calls to create a completion return our mock
+    // This is handled by the inherited Chat.Completions mock in openai
+    const response = await client.chat.completions.create({
+      model: 'gpt-4',
+      messages: [{ role: 'user', content: 'Hello' }],
+      posthog_distinct_id: 'test-id',
+      posthog_properties: { foo: 'bar' },
+    })
+
+    expect(response).toEqual(mockOpenAiChatResponse)
+    // 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('openai')
+    expect(properties['$ai_model']).toBe('gpt-4')
+    expect(properties['$ai_input']).toEqual([{ role: 'user', content: 'Hello' }])
+    expect(properties['$ai_output_choices']).toEqual([{ role: 'assistant', content: 'Hello from OpenAI!' }])
+    expect(properties['$ai_input_tokens']).toBe(20)
+    expect(properties['$ai_output_tokens']).toBe(10)
+    expect(properties['$ai_http_status']).toBe(200)
+    expect(properties['foo']).toBe('bar')
+    expect(typeof properties['$ai_latency']).toBe('number')
+  })
+
+  conditionalTest('embeddings', async () => {
+    // Since embeddings calls are not implemented in the snippet by default,
+    // we'll demonstrate how you *would* do it if WrappedEmbeddings is used.
+    // Let's override the internal embeddings to return our mock.
+    const mockEmbeddingsCreate = jest.fn().mockResolvedValue(mockOpenAiEmbeddingResponse)
+    ;(client as any).embeddings = {
+      create: mockEmbeddingsCreate,
+    }
+
+    const response = await (client as any).embeddings.create({
+      model: 'text-embedding-3-small',
+      input: 'Hello world',
+      posthog_distinct_id: 'test-id',
+      posthog_properties: { foo: 'bar' },
+    })
+
+    expect(response).toEqual(mockOpenAiEmbeddingResponse)
+    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('openai')
+    expect(properties['$ai_model']).toBe('text-embedding-3-small')
+    expect(properties['$ai_input']).toBe('Hello world')
+    expect(properties['$ai_input_tokens']).toBe(10)
+    expect(properties['$ai_http_status']).toBe(200)
+    expect(properties['foo']).toBe('bar')
+    expect(typeof properties['$ai_latency']).toBe('number')
+  })
+
+  conditionalTest('groups', async () => {
+    await client.chat.completions.create({
+      model: 'gpt-4',
+      messages: [{ role: 'user', content: 'Hello' }],
+      posthog_distinct_id: 'test-id',
+      posthog_groups: { company: 'test_company' },
+    })
+    expect(mockPostHogClient.capture).toHaveBeenCalledTimes(1)
+    const [captureArgs] = (mockPostHogClient.capture as jest.Mock).mock.calls
+    const { groups } = captureArgs[0]
+    expect(groups).toEqual({ company: 'test_company' })
+  })
+
+  conditionalTest('privacy mode local', async () => {
+    await client.chat.completions.create({
+      model: 'gpt-4',
+      messages: [{ role: 'user', content: 'Hello' }],
+      posthog_distinct_id: 'test-id',
+      posthog_privacy_mode: 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('privacy mode global', async () => {
+    // override mock to appear globally in privacy mode
+    ;(mockPostHogClient as any).privacy_mode = true
+
+    await client.chat.completions.create({
+      model: 'gpt-4',
+      messages: [{ role: 'user', content: 'Hello' }],
+      posthog_distinct_id: 'test-id',
+      // we attempt to override locally, but it should still be null if global is true
+      posthog_privacy_mode: false,
+    })
+
+    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('core model params', async () => {
+    mockOpenAiChatResponse.usage = {
+      prompt_tokens: 20,
+      completion_tokens: 10,
+    }
+
+    await client.chat.completions.create({
+      model: 'gpt-4',
+      // using openai-like params
+      temperature: 0.5,
+      max_completion_tokens: 100,
+      stream: false,
+      messages: [{ role: 'user', content: 'Hello' }],
+      posthog_distinct_id: 'test-id',
+      posthog_properties: { foo: 'bar' },
+    })
+
+    expect(mockPostHogClient.capture).toHaveBeenCalledTimes(1)
+    const [captureArgs] = (mockPostHogClient.capture as jest.Mock).mock.calls
+    const { properties } = captureArgs[0]
+
+    expect(properties['$ai_model_parameters']).toEqual({
+      temperature: 0.5,
+      max_completion_tokens: 100,
+      stream: false,
+    })
+    expect(properties['$ai_temperature']).toBe(0.5)
+    expect(properties['$ai_max_tokens']).toBe(100)
+    expect(properties['$ai_stream']).toBe(false)
+    expect(properties['foo']).toBe('bar')
+  })
+})

package/tsconfig.json

@@ -0,0 +1,10 @@
+{
+  "extends": "../tsconfig.json",
+  "compilerOptions": {
+    "types": ["node"],
+    "typeRoots": ["./node_modules/@types", "../node_modules/@types"],
+    "moduleResolution": "node",
+    "skipLibCheck": true,
+    "allowSyntheticDefaultImports": true
+  }
+}