@posthog/ai 5.0.1 → 5.2.0

package/src/openai/azure.ts

@@ -2,14 +2,18 @@ import OpenAIOrignal, { AzureOpenAI } from 'openai'
 import { PostHog } from 'posthog-node'
 import { v4 as uuidv4 } from 'uuid'
 import { formatResponseOpenAI, MonitoringParams, sendEventToPosthog } from '../utils'
+import type { APIPromise } from 'openai'
+import type { Stream } from 'openai/streaming'
+import type { ParsedResponse } from 'openai/resources/responses/responses'
 
 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'
+type ResponsesCreateParamsBase = OpenAIOrignal.Responses.ResponseCreateParams
+type ResponsesCreateParamsNonStreaming = OpenAIOrignal.Responses.ResponseCreateParamsNonStreaming
+type ResponsesCreateParamsStreaming = OpenAIOrignal.Responses.ResponseCreateParamsStreaming
 
 interface MonitoringOpenAIConfig {
   apiKey: string
@@ -17,6 +21,8 @@ interface MonitoringOpenAIConfig {
   baseURL?: string
 }
 
+type RequestOptions = Record<string, any>
+
 export class PostHogAzureOpenAI extends AzureOpenAI {
   private readonly phClient: PostHog
   public chat: WrappedChat
@@ -82,32 +88,30 @@ export class WrappedCompletions extends AzureOpenAI.Chat.Completions {
 
     const traceId = posthogTraceId ?? uuidv4()
     const startTime = Date.now()
+
     const parentPromise = super.create(openAIParams, options)
 
     if (openAIParams.stream) {
       return parentPromise.then((value) => {
-        let accumulatedContent = ''
-        let usage: {
-          inputTokens: number
-          outputTokens: number
-          reasoningTokens?: number
-          cacheReadInputTokens?: number
-        } = {
-          inputTokens: 0,
-          outputTokens: 0,
-        }
-        let model = openAIParams.model
         if ('tee' in value) {
           const [stream1, stream2] = value.tee()
           ;(async () => {
             try {
+              let accumulatedContent = ''
+              let usage: {
+                inputTokens?: number
+                outputTokens?: number
+                reasoningTokens?: number
+                cacheReadInputTokens?: number
+              } = {
+                inputTokens: 0,
+                outputTokens: 0,
+              }
+
               for await (const chunk of stream1) {
                 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,
@@ -116,12 +120,13 @@ export class WrappedCompletions extends AzureOpenAI.Chat.Completions {
                   }
                 }
               }
+
               const latency = (Date.now() - startTime) / 1000
               await sendEventToPosthog({
                 client: this.phClient,
-                distinctId: posthogDistinctId ?? traceId,
+                distinctId: posthogDistinctId,
                 traceId,
-                model,
+                model: openAIParams.model,
                 provider: 'azure',
                 input: openAIParams.messages,
                 output: [{ content: accumulatedContent, role: 'assistant' }],
@@ -133,23 +138,19 @@ export class WrappedCompletions extends AzureOpenAI.Chat.Completions {
                 captureImmediate: posthogCaptureImmediate,
               })
             } catch (error: any) {
-              // error handling
               await sendEventToPosthog({
                 client: this.phClient,
-                distinctId: posthogDistinctId ?? traceId,
+                distinctId: posthogDistinctId,
                 traceId,
-                model,
+                model: openAIParams.model,
                 provider: 'azure',
                 input: openAIParams.messages,
-                output: JSON.stringify(error),
+                output: [],
                 latency: 0,
                 baseURL: (this as any).baseURL ?? '',
                 params: body,
                 httpStatus: error?.status ? error.status : 500,
-                usage: {
-                  inputTokens: 0,
-                  outputTokens: 0,
-                },
+                usage: { inputTokens: 0, outputTokens: 0 },
                 isError: true,
                 error: JSON.stringify(error),
                 captureImmediate: posthogCaptureImmediate,
@@ -167,15 +168,11 @@ export class WrappedCompletions extends AzureOpenAI.Chat.Completions {
         async (result) => {
           if ('choices' in result) {
             const latency = (Date.now() - startTime) / 1000
-            let model = openAIParams.model
-            if (result.model != model) {
-              model = result.model
-            }
             await sendEventToPosthog({
               client: this.phClient,
-              distinctId: posthogDistinctId ?? traceId,
+              distinctId: posthogDistinctId,
               traceId,
-              model,
+              model: openAIParams.model,
               provider: 'azure',
               input: openAIParams.messages,
               output: formatResponseOpenAI(result),
@@ -197,7 +194,7 @@ export class WrappedCompletions extends AzureOpenAI.Chat.Completions {
         async (error: any) => {
           await sendEventToPosthog({
             client: this.phClient,
-            distinctId: posthogDistinctId ?? traceId,
+            distinctId: posthogDistinctId,
             traceId,
             model: openAIParams.model,
             provider: 'azure',
@@ -224,4 +221,261 @@ export class WrappedCompletions extends AzureOpenAI.Chat.Completions {
   }
 }
 
+export class WrappedResponses extends AzureOpenAI.Responses {
+  private readonly phClient: PostHog
+
+  constructor(client: AzureOpenAI, phClient: PostHog) {
+    super(client)
+    this.phClient = phClient
+  }
+
+  // --- Overload #1: Non-streaming
+  public create(
+    body: ResponsesCreateParamsNonStreaming & MonitoringParams,
+    options?: RequestOptions
+  ): APIPromise<OpenAIOrignal.Responses.Response>
+
+  // --- Overload #2: Streaming
+  public create(
+    body: ResponsesCreateParamsStreaming & MonitoringParams,
+    options?: RequestOptions
+  ): APIPromise<Stream<OpenAIOrignal.Responses.ResponseStreamEvent>>
+
+  // --- Overload #3: Generic base
+  public create(
+    body: ResponsesCreateParamsBase & MonitoringParams,
+    options?: RequestOptions
+  ): APIPromise<OpenAIOrignal.Responses.Response | Stream<OpenAIOrignal.Responses.ResponseStreamEvent>>
+
+  // --- Implementation Signature
+  public create(
+    body: ResponsesCreateParamsBase & MonitoringParams,
+    options?: RequestOptions
+  ): APIPromise<OpenAIOrignal.Responses.Response | Stream<OpenAIOrignal.Responses.ResponseStreamEvent>> {
+    const {
+      posthogDistinctId,
+      posthogTraceId,
+      posthogProperties,
+      // eslint-disable-next-line @typescript-eslint/no-unused-vars
+      posthogPrivacyMode = false,
+      posthogGroups,
+      posthogCaptureImmediate,
+      ...openAIParams
+    } = body
+
+    const traceId = posthogTraceId ?? uuidv4()
+    const startTime = Date.now()
+
+    const parentPromise = super.create(openAIParams, options)
+
+    if (openAIParams.stream) {
+      return parentPromise.then((value) => {
+        if ('tee' in value && typeof (value as any).tee === 'function') {
+          const [stream1, stream2] = (value as any).tee()
+          ;(async () => {
+            try {
+              let finalContent: any[] = []
+              let usage: {
+                inputTokens?: number
+                outputTokens?: number
+                reasoningTokens?: number
+                cacheReadInputTokens?: number
+              } = {
+                inputTokens: 0,
+                outputTokens: 0,
+              }
+
+              for await (const chunk of stream1) {
+                if (
+                  chunk.type === 'response.completed' &&
+                  'response' in chunk &&
+                  chunk.response?.output &&
+                  chunk.response.output.length > 0
+                ) {
+                  finalContent = chunk.response.output
+                }
+                if ('usage' in chunk && chunk.usage) {
+                  usage = {
+                    inputTokens: chunk.usage.input_tokens ?? 0,
+                    outputTokens: chunk.usage.output_tokens ?? 0,
+                    reasoningTokens: chunk.usage.output_tokens_details?.reasoning_tokens ?? 0,
+                    cacheReadInputTokens: chunk.usage.input_tokens_details?.cached_tokens ?? 0,
+                  }
+                }
+              }
+
+              const latency = (Date.now() - startTime) / 1000
+              await sendEventToPosthog({
+                client: this.phClient,
+                distinctId: posthogDistinctId,
+                traceId,
+                model: openAIParams.model,
+                provider: 'azure',
+                input: openAIParams.input,
+                output: finalContent,
+                latency,
+                baseURL: (this as any).baseURL ?? '',
+                params: body,
+                httpStatus: 200,
+                usage,
+                captureImmediate: posthogCaptureImmediate,
+              })
+            } catch (error: any) {
+              await sendEventToPosthog({
+                client: this.phClient,
+                distinctId: posthogDistinctId,
+                traceId,
+                model: openAIParams.model,
+                provider: 'azure',
+                input: openAIParams.input,
+                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),
+                captureImmediate: posthogCaptureImmediate,
+              })
+            }
+          })()
+
+          return stream2
+        }
+        return value
+      }) as APIPromise<Stream<OpenAIOrignal.Responses.ResponseStreamEvent>>
+    } else {
+      const wrappedPromise = parentPromise.then(
+        async (result) => {
+          if ('output' in result) {
+            const latency = (Date.now() - startTime) / 1000
+            await sendEventToPosthog({
+              client: this.phClient,
+              distinctId: posthogDistinctId,
+              traceId,
+              model: openAIParams.model,
+              provider: 'azure',
+              input: openAIParams.input,
+              output: result.output,
+              latency,
+              baseURL: (this as any).baseURL ?? '',
+              params: body,
+              httpStatus: 200,
+              usage: {
+                inputTokens: result.usage?.input_tokens ?? 0,
+                outputTokens: result.usage?.output_tokens ?? 0,
+                reasoningTokens: result.usage?.output_tokens_details?.reasoning_tokens ?? 0,
+                cacheReadInputTokens: result.usage?.input_tokens_details?.cached_tokens ?? 0,
+              },
+              captureImmediate: posthogCaptureImmediate,
+            })
+          }
+          return result
+        },
+        async (error: any) => {
+          await sendEventToPosthog({
+            client: this.phClient,
+            distinctId: posthogDistinctId,
+            traceId,
+            model: openAIParams.model,
+            provider: 'azure',
+            input: openAIParams.input,
+            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),
+            captureImmediate: posthogCaptureImmediate,
+          })
+          throw error
+        }
+      ) as APIPromise<OpenAIOrignal.Responses.Response>
+
+      return wrappedPromise
+    }
+  }
+
+  public parse<Params extends OpenAIOrignal.Responses.ResponseCreateParams, ParsedT = any>(
+    body: Params & MonitoringParams,
+    options?: RequestOptions
+  ): APIPromise<ParsedResponse<ParsedT>> {
+    const {
+      posthogDistinctId,
+      posthogTraceId,
+      posthogProperties,
+      // eslint-disable-next-line @typescript-eslint/no-unused-vars
+      posthogPrivacyMode = false,
+      posthogGroups,
+      posthogCaptureImmediate,
+      ...openAIParams
+    } = body
+
+    const traceId = posthogTraceId ?? uuidv4()
+    const startTime = Date.now()
+
+    const parentPromise = super.parse(openAIParams, options)
+
+    const wrappedPromise = parentPromise.then(
+      async (result) => {
+        const latency = (Date.now() - startTime) / 1000
+        await sendEventToPosthog({
+          client: this.phClient,
+          distinctId: posthogDistinctId,
+          traceId,
+          model: openAIParams.model,
+          provider: 'azure',
+          input: openAIParams.input,
+          output: result.output,
+          latency,
+          baseURL: (this as any).baseURL ?? '',
+          params: body,
+          httpStatus: 200,
+          usage: {
+            inputTokens: result.usage?.input_tokens ?? 0,
+            outputTokens: result.usage?.output_tokens ?? 0,
+            reasoningTokens: result.usage?.output_tokens_details?.reasoning_tokens ?? 0,
+            cacheReadInputTokens: result.usage?.input_tokens_details?.cached_tokens ?? 0,
+          },
+          captureImmediate: posthogCaptureImmediate,
+        })
+        return result
+      },
+      async (error: any) => {
+        await sendEventToPosthog({
+          client: this.phClient,
+          distinctId: posthogDistinctId,
+          traceId,
+          model: openAIParams.model,
+          provider: 'azure',
+          input: openAIParams.input,
+          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),
+          captureImmediate: posthogCaptureImmediate,
+        })
+        throw error
+      }
+    )
+
+    return wrappedPromise as APIPromise<ParsedResponse<ParsedT>>
+  }
+}
+
 export default PostHogAzureOpenAI
+
+export { PostHogAzureOpenAI as OpenAI }