@posthog/ai 2.1.2 → 2.3.0

package/src/openai/azure.ts

@@ -0,0 +1,214 @@
+import OpenAIOrignal, { AzureOpenAI } from 'openai'
+import { PostHog } from 'posthog-node'
+import { v4 as uuidv4 } from 'uuid'
+import { PassThrough } from 'stream'
+import { mergeSystemPrompt, 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
+  public chat: WrappedChat
+
+  constructor(config: MonitoringOpenAIConfig) {
+    const { posthog, ...openAIConfig } = config
+    super(openAIConfig)
+    this.phClient = posthog
+    this.chat = new WrappedChat(this, this.phClient)
+  }
+}
+
+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 {
+      posthogDistinctId,
+      posthogTraceId,
+      posthogProperties,
+      posthogPrivacyMode = false,
+      posthogGroups,
+      ...openAIParams
+    } = body
+
+    const traceId = posthogTraceId ?? 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: { inputTokens: number; outputTokens: number } = {
+          inputTokens: 0,
+          outputTokens: 0,
+        }
+        let model = openAIParams.model
+        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) {
+                  if (chunk.model != model) {
+                    model = chunk.model
+                  }
+                  usage = {
+                    inputTokens: chunk.usage.prompt_tokens ?? 0,
+                    outputTokens: chunk.usage.completion_tokens ?? 0,
+                  }
+                }
+                passThroughStream.write(chunk)
+              }
+              const latency = (Date.now() - startTime) / 1000
+              sendEventToPosthog({
+                client: this.phClient,
+                distinctId: posthogDistinctId ?? traceId,
+                traceId,
+                model,
+                provider: 'azure',
+                input: mergeSystemPrompt(openAIParams, 'azure'),
+                output: [{ content: accumulatedContent, role: 'assistant' }],
+                latency,
+                baseURL: (this as any).baseURL ?? '',
+                params: body,
+                httpStatus: 200,
+                usage,
+              })
+              passThroughStream.end()
+            } catch (error: any) {
+              // error handling
+              sendEventToPosthog({
+                client: this.phClient,
+                distinctId: posthogDistinctId ?? traceId,
+                traceId,
+                model,
+                provider: 'azure',
+                input: mergeSystemPrompt(openAIParams, 'azure'),
+                output: JSON.stringify(error),
+                latency: 0,
+                baseURL: (this as any).baseURL ?? '',
+                params: body,
+                httpStatus: error?.status ? error.status : 500,
+                usage: {
+                  inputTokens: 0,
+                  outputTokens: 0,
+                },
+                isError: true,
+                error: JSON.stringify(error),
+              })
+              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
+            let model = openAIParams.model
+            if (result.model != model) {
+              model = result.model
+            }
+            sendEventToPosthog({
+              client: this.phClient,
+              distinctId: posthogDistinctId ?? traceId,
+              traceId,
+              model,
+              provider: '1234',
+              input: mergeSystemPrompt(openAIParams, 'azure'),
+              output: [{ content: result.choices[0].message.content, role: 'assistant' }],
+              latency,
+              baseURL: (this as any).baseURL ?? '',
+              params: body,
+              httpStatus: 200,
+              usage: {
+                inputTokens: result.usage?.prompt_tokens ?? 0,
+                outputTokens: result.usage?.completion_tokens ?? 0,
+              },
+            })
+          }
+          return result
+        },
+        (error: any) => {
+          sendEventToPosthog({
+            client: this.phClient,
+            distinctId: posthogDistinctId ?? traceId,
+            traceId,
+            model: openAIParams.model,
+            provider: 'azure',
+            input: mergeSystemPrompt(openAIParams, 'azure'),
+            output: [],
+            latency: 0,
+            baseURL: (this as any).baseURL ?? '',
+            params: body,
+            httpStatus: error?.status ? error.status : 500,
+            usage: {
+              inputTokens: 0,
+              outputTokens: 0,
+            },
+            isError: true,
+            error: JSON.stringify(error),
+          })
+          throw error
+        }
+      ) as APIPromise<ChatCompletion>
+
+      return wrappedPromise
+    }
+  }
+}
+
+export default PostHogAzureOpenAI

package/src/openai/index.ts

@@ -123,7 +123,7 @@ export class WrappedCompletions extends OpenAIOrignal.Chat.Completions {
                 usage,
               })
               passThroughStream.end()
-            } catch (error) {
+            } catch (error: any) {
               // error handling
               sendEventToPosthog({
                 client: this.phClient,
@@ -136,11 +136,13 @@ export class WrappedCompletions extends OpenAIOrignal.Chat.Completions {
                 latency: 0,
                 baseURL: (this as any).baseURL ?? '',
                 params: body,
-                httpStatus: 500,
+                httpStatus: error?.status ? error.status : 500,
                 usage: {
                   inputTokens: 0,
                   outputTokens: 0,
                 },
+                isError: true,
+                error: JSON.stringify(error),
               })
               passThroughStream.emit('error', error)
             }
@@ -173,7 +175,7 @@ export class WrappedCompletions extends OpenAIOrignal.Chat.Completions {
           }
           return result
         },
-        (error) => {
+        (error: any) => {
           sendEventToPosthog({
             client: this.phClient,
             distinctId: posthogDistinctId ?? traceId,
@@ -185,11 +187,13 @@ export class WrappedCompletions extends OpenAIOrignal.Chat.Completions {
             latency: 0,
             baseURL: (this as any).baseURL ?? '',
             params: body,
-            httpStatus: 500,
+            httpStatus: error?.status ? error.status : 500,
             usage: {
               inputTokens: 0,
               outputTokens: 0,
             },
+            isError: true,
+            error: JSON.stringify(error),
           })
           throw error
         }

package/src/utils.ts

@@ -77,15 +77,15 @@ export const formatResponseOpenAI = (response: any): Array<{ role: string; conte
 }
 
 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
+  if (provider == 'anthropic') {
+    const messages = params.messages || []
+    if (!(params as any).system) {
+      return messages
+    }
+    const systemMessage = (params as any).system
+    return [{ role: 'system', content: systemMessage }, ...messages]
   }
-  const systemMessage = (params as any).system
-  return [{ role: 'system', content: systemMessage }, ...messages]
+  return params.messages
 }
 
 export const withPrivacyMode = (client: PostHog, privacyMode: boolean, input: any): any => {
@@ -105,6 +105,8 @@ export type SendEventToPosthogParams = {
   httpStatus: number
   usage?: { inputTokens?: number; outputTokens?: number }
   params: ChatCompletionCreateParamsBase & MonitoringParams
+  isError?: boolean
+  error?: string
 }
 
 export const sendEventToPosthog = ({
@@ -120,8 +122,17 @@ export const sendEventToPosthog = ({
   params,
   httpStatus = 200,
   usage = {},
+  isError = false,
+  error,
 }: SendEventToPosthogParams): void => {
   if (client.capture) {
+    let errorData = {}
+    if (isError) {
+      errorData = {
+        $ai_is_error: true,
+        $ai_error: error,
+      }
+    }
     client.capture({
       distinctId: distinctId ?? traceId,
       event: '$ai_generation',
@@ -139,6 +150,7 @@ export const sendEventToPosthog = ({
         $ai_base_url: baseURL,
         ...params.posthogProperties,
         ...(distinctId ? {} : { $process_person_profile: false }),
+        ...errorData,
       },
       groups: params.posthogGroups,
     })

package/src/vercel/middleware.ts

@@ -15,6 +15,8 @@ interface ClientOptions {
   posthogProperties?: Record<string, any>
   posthogPrivacyMode?: boolean
   posthogGroups?: Record<string, any>
+  posthogModelOverride?: string
+  posthogProviderOverride?: string
 }
 
 interface CreateInstrumentationMiddlewareOptions {
@@ -23,6 +25,8 @@ interface CreateInstrumentationMiddlewareOptions {
   posthogProperties?: Record<string, any>
   posthogPrivacyMode?: boolean
   posthogGroups?: Record<string, any>
+  posthogModelOverride?: string
+  posthogProviderOverride?: string
 }
 
 interface PostHogInput {
@@ -80,14 +84,15 @@ export const createInstrumentationMiddleware = (
         const result = await doGenerate()
         const latency = (Date.now() - startTime) / 1000
 
-        const modelId = result.response?.modelId ? result.response.modelId : model.modelId
+        const modelId = options.posthogModelOverride ?? (result.response?.modelId ? result.response.modelId : model.modelId)
+        const provider = options.posthogProviderOverride ?? model.provider
 
         sendEventToPosthog({
           client: phClient,
           distinctId: options.posthogDistinctId,
           traceId: options.posthogTraceId,
           model: modelId,
-          provider: model.provider,
+          provider: provider,
           input: options.posthogPrivacyMode ? '' : mapVercelPrompt(params.prompt),
           output: [{ content: result.text, role: 'assistant' }],
           latency,
@@ -101,7 +106,7 @@ export const createInstrumentationMiddleware = (
         })
 
         return result
-      } catch (error) {
+      } catch (error: any) {
         const modelId = model.modelId
         sendEventToPosthog({
           client: phClient,
@@ -114,11 +119,13 @@ export const createInstrumentationMiddleware = (
           latency: 0,
           baseURL: '',
           params: mergedParams as any,
-          httpStatus: 500,
+          httpStatus: error?.status ? error.status : 500,
           usage: {
             inputTokens: 0,
             outputTokens: 0,
           },
+          isError: true,
+          error: JSON.stringify(error),
         })
         throw error
       }
@@ -132,6 +139,9 @@ export const createInstrumentationMiddleware = (
         ...options,
         ...mapVercelParams(params),
       }
+
+      const modelId = options.posthogModelOverride ?? model.modelId
+      const provider = options.posthogProviderOverride ?? model.provider
       try {
         const { stream, ...rest } = await doStream()
 
@@ -155,8 +165,8 @@ export const createInstrumentationMiddleware = (
               client: phClient,
               distinctId: options.posthogDistinctId,
               traceId: options.posthogTraceId,
-              model: model.modelId,
-              provider: model.provider,
+              model: modelId,
+              provider: provider,
               input: options.posthogPrivacyMode ? '' : mapVercelPrompt(params.prompt),
               output: [{ content: generatedText, role: 'assistant' }],
               latency,
@@ -172,23 +182,25 @@ export const createInstrumentationMiddleware = (
           stream: stream.pipeThrough(transformStream),
           ...rest,
         }
-      } catch (error) {
+      } catch (error: any) {
         sendEventToPosthog({
           client: phClient,
           distinctId: options.posthogDistinctId,
           traceId: options.posthogTraceId,
-          model: model.modelId,
-          provider: model.provider,
+          model: modelId,
+          provider: provider,
           input: options.posthogPrivacyMode ? '' : mapVercelPrompt(params.prompt),
           output: [],
           latency: 0,
           baseURL: '',
           params: mergedParams as any,
-          httpStatus: 500,
+          httpStatus: error?.status ? error.status : 500,
           usage: {
             inputTokens: 0,
             outputTokens: 0,
           },
+          isError: true,
+          error: JSON.stringify(error),
         })
         throw error
       }