feat(llma): add $ai_stop_reason extraction for all providers

.changeset/add-ai-stop-reason.md

@@ -0,0 +1,5 @@
+---
+'@posthog/ai': minor
+---
+
+Add `$ai_stop_reason` property capturing the LLM's reason for stopping generation across all providers

packages/ai/src/anthropic/index.ts

@@ -79,6 +79,7 @@ export class WrappedMessages extends AnthropicOriginal.Messages {
         const toolsInProgress: Map<string, ToolInProgress> = new Map()
         let currentTextBlock: FormattedTextContent | null = null
         let firstTokenTime: number | undefined
+        let stopReason: string | undefined
 
         const usage: {
           inputTokens: number
@@ -202,6 +203,13 @@ export class WrappedMessages extends AnthropicOriginal.Messages {
                     usage.webSearchCount = chunk.usage.server_tool_use.web_search_requests
                   }
                 }
+
+                if (chunk.type === 'message_delta' && 'delta' in chunk) {
+                  const delta = chunk.delta
+                  if ('stop_reason' in delta && typeof delta.stop_reason === 'string' && delta.stop_reason) {
+                    stopReason = delta.stop_reason
+                  }
+                }
               }
               usage.rawUsage = lastRawUsage
 
@@ -239,6 +247,7 @@ export class WrappedMessages extends AnthropicOriginal.Messages {
                 params: body,
                 httpStatus: 200,
                 usage,
+                stopReason,
                 tools: availableTools,
               })
             } catch (error: unknown) {
@@ -294,6 +303,7 @@ export class WrappedMessages extends AnthropicOriginal.Messages {
                 webSearchCount: result.usage.server_tool_use?.web_search_requests ?? 0,
                 rawUsage: result.usage,
               },
+              stopReason: result.stop_reason ?? undefined,
               tools: availableTools,
             })
           }

packages/ai/src/gemini/index.ts

@@ -4,6 +4,8 @@ import {
   GenerateContentParameters,
   Part,
   GenerateContentResponseUsageMetadata,
+  EmbedContentParameters,
+  EmbedContentResponse,
 } from '@google/genai'
 import type { GoogleGenAIOptions } from '@google/genai'
 import { PostHog } from 'posthog-node'
@@ -15,6 +17,8 @@ import {
   extractPosthogParams,
   toContentString,
   sendEventWithErrorToPosthog,
+  AIEvent,
+  withPrivacyMode,
 } from '../utils'
 import { sanitizeGemini } from '../sanitization'
 import type { TokenUsage, FormattedContent, FormattedContentItem, FormattedMessage } from '../types'
@@ -57,6 +61,7 @@ export class WrappedModels {
       const availableTools = extractAvailableToolCalls('gemini', geminiParams)
 
       const metadata = response.usageMetadata
+      const finishReason = response.candidates?.[0]?.finishReason
       await sendEventToPosthog({
         client: this.phClient,
         ...posthogParams,
@@ -78,6 +83,7 @@ export class WrappedModels {
           webSearchCount: calculateGoogleWebSearchCount(response),
           rawUsage: metadata,
         },
+        stopReason: finishReason ?? undefined,
         tools: availableTools,
       })
 
@@ -111,6 +117,7 @@ export class WrappedModels {
     const startTime = Date.now()
     const accumulatedContent: FormattedContent = []
     let firstTokenTime: number | undefined
+    let stopReason: string | undefined
     let usage: TokenUsage = {
       inputTokens: 0,
       outputTokens: 0,
@@ -149,6 +156,11 @@ export class WrappedModels {
           }
         }
 
+        // Track finish reason from candidates
+        if (chunk.candidates?.[0]?.finishReason) {
+          stopReason = chunk.candidates[0].finishReason
+        }
+
         // Handle function calls from candidates
         if (chunk.candidates && Array.isArray(chunk.candidates)) {
           for (const candidate of chunk.candidates) {
@@ -217,6 +229,7 @@ export class WrappedModels {
           webSearchCount: usage.webSearchCount,
           rawUsage: usage.rawUsage,
         },
+        stopReason,
         tools: availableTools,
       })
     } catch (error: unknown) {
@@ -241,6 +254,57 @@ export class WrappedModels {
     }
   }
 
+  public async embedContent(params: EmbedContentParameters & MonitoringParams): Promise<EmbedContentResponse> {
+    const { providerParams: geminiParams, posthogParams } = extractPosthogParams(params)
+    const startTime = Date.now()
+
+    try {
+      const response = await this.client.models.embedContent(geminiParams as EmbedContentParameters)
+      const latency = (Date.now() - startTime) / 1000
+
+      const tokenCount =
+        response.embeddings?.reduce((sum, embedding) => sum + (embedding.statistics?.tokenCount ?? 0), 0) ?? 0
+
+      await sendEventToPosthog({
+        client: this.phClient,
+        ...posthogParams,
+        eventType: AIEvent.Embedding,
+        model: geminiParams.model,
+        provider: 'gemini',
+        input: withPrivacyMode(this.phClient, posthogParams.privacyMode, geminiParams.contents),
+        output: null,
+        latency,
+        baseURL: 'https://generativelanguage.googleapis.com',
+        params: params as EmbedContentParameters & MonitoringParams,
+        httpStatus: 200,
+        usage: {
+          inputTokens: tokenCount,
+        },
+      })
+
+      return response
+    } catch (error: unknown) {
+      const latency = (Date.now() - startTime) / 1000
+      const enrichedError = await sendEventWithErrorToPosthog({
+        client: this.phClient,
+        ...posthogParams,
+        eventType: AIEvent.Embedding,
+        model: geminiParams.model,
+        provider: 'gemini',
+        input: withPrivacyMode(this.phClient, posthogParams.privacyMode, geminiParams.contents),
+        output: null,
+        latency,
+        baseURL: 'https://generativelanguage.googleapis.com',
+        params: params as EmbedContentParameters & MonitoringParams,
+        usage: {
+          inputTokens: 0,
+        },
+        error: error,
+      })
+      throw enrichedError
+    }
+  }
+
   private formatPartsAsContentBlocks(parts: unknown[]): FormattedContent {
     const blocks: FormattedContent = []
 

packages/ai/src/langchain/callbacks.ts

@@ -458,6 +458,12 @@ export class LangChainCallbackHandler extends BaseCallbackHandler {
         eventProperties['$ai_web_search_count'] = additionalTokenData.webSearchCount
       }
 
+      // Extract stop reason from generation info
+      const stopReason = this._extractStopReason(output)
+      if (stopReason) {
+        eventProperties['$ai_stop_reason'] = stopReason
+      }
+
       // Handle generations/completions
       let completions
       if (output.generations && Array.isArray(output.generations)) {
@@ -588,6 +594,39 @@ export class LangChainCallbackHandler extends BaseCallbackHandler {
     return sanitizeLangChain(messageDict) as Record<string, any>
   }
 
+  private _extractStopReason(output: LLMResult): string | undefined {
+    if (!output.generations || !Array.isArray(output.generations)) {
+      return undefined
+    }
+    const lastGeneration = output.generations[output.generations.length - 1]
+    if (!Array.isArray(lastGeneration) || lastGeneration.length === 0) {
+      return undefined
+    }
+    const gen = lastGeneration[0]
+
+    // Check generationInfo for finish_reason (OpenAI format)
+    if (gen.generationInfo?.finish_reason) {
+      return String(gen.generationInfo.finish_reason)
+    }
+
+    // Check generationInfo for response_metadata.stop_reason (Anthropic format)
+    if (gen.generationInfo?.response_metadata?.stop_reason) {
+      return String(gen.generationInfo.response_metadata.stop_reason)
+    }
+
+    // Check message response_metadata for finish_reason (common LangChain format)
+    if (gen.generationInfo?.response_metadata?.finish_reason) {
+      return String(gen.generationInfo.response_metadata.finish_reason)
+    }
+
+    // Check for stop_reason directly in generationInfo
+    if (gen.generationInfo?.stop_reason) {
+      return String(gen.generationInfo.stop_reason)
+    }
+
+    return undefined
+  }
+
   private _parseUsageModel(usage: any, provider?: string, model?: string): [number, number, Record<string, any>] {
     const conversionList: Array<[string, 'input' | 'output']> = [
       ['promptTokens', 'input'],

packages/ai/src/openai/index.ts

@@ -121,6 +121,7 @@ export class WrappedCompletions extends Completions {
               let accumulatedContent = ''
               let modelFromResponse: string | undefined
               let firstTokenTime: number | undefined
+              let stopReason: string | undefined
               let usage: {
                 inputTokens?: number
                 outputTokens?: number
@@ -152,6 +153,10 @@ export class WrappedCompletions extends Completions {
 
                 const choice = chunk?.choices?.[0]
 
+                if (choice?.finish_reason) {
+                  stopReason = choice.finish_reason
+                }
+
                 const chunkWebSearchCount = calculateWebSearchCount(chunk)
                 if (chunkWebSearchCount > 0 && chunkWebSearchCount > (usage.webSearchCount ?? 0)) {
                   usage.webSearchCount = chunkWebSearchCount
@@ -273,6 +278,7 @@ export class WrappedCompletions extends Completions {
                   webSearchCount: usage.webSearchCount,
                   rawUsage: rawUsageData,
                 },
+                stopReason,
                 tools: availableTools,
               })
             } catch (error: unknown) {
@@ -324,6 +330,7 @@ export class WrappedCompletions extends Completions {
                 webSearchCount: calculateWebSearchCount(result),
                 rawUsage: result.usage,
               },
+              stopReason: result.choices[0]?.finish_reason ?? undefined,
               tools: availableTools,
             })
           }
@@ -408,6 +415,7 @@ export class WrappedResponses extends Responses {
               let finalContent: unknown[] = []
               let modelFromResponse: string | undefined
               let firstTokenTime: number | undefined
+              let stopReason: string | undefined
               let usage: {
                 inputTokens?: number
                 outputTokens?: number
@@ -446,6 +454,9 @@ export class WrappedResponses extends Responses {
                   chunk.response.output.length > 0
                 ) {
                   finalContent = chunk.response.output
+                  if (chunk.response.status) {
+                    stopReason = chunk.response.status
+                  }
                 }
                 if ('response' in chunk && chunk.response?.usage) {
                   rawUsageData = chunk.response.usage
@@ -485,6 +496,7 @@ export class WrappedResponses extends Responses {
                   webSearchCount: usage.webSearchCount,
                   rawUsage: rawUsageData,
                 },
+                stopReason,
                 tools: availableTools,
               })
             } catch (error: unknown) {
@@ -538,6 +550,7 @@ export class WrappedResponses extends Responses {
                 webSearchCount: calculateWebSearchCount(result),
                 rawUsage: result.usage,
               },
+              stopReason: result.status ?? undefined,
               tools: availableTools,
             })
           }
@@ -610,6 +623,7 @@ export class WrappedResponses extends Responses {
               cacheReadInputTokens: result.usage?.input_tokens_details?.cached_tokens ?? 0,
               rawUsage: result.usage,
             },
+            stopReason: result.status ?? undefined,
           })
           return result
         },

packages/ai/src/utils.ts

@@ -582,6 +582,7 @@ export type SendEventToPosthogParams = {
     MonitoringParams
   error?: unknown
   exceptionId?: string
+  stopReason?: string
   tools?: ChatCompletionTool[] | AnthropicTool[] | GeminiTool[] | null
   captureImmediate?: boolean
 }
@@ -688,6 +689,7 @@ export const sendEventToPosthog = async ({
   usage = {},
   error,
   exceptionId,
+  stopReason,
   tools,
   captureImmediate = false,
 }: SendEventToPosthogParams): Promise<void> => {
@@ -745,6 +747,7 @@ export const sendEventToPosthog = async ({
     ...params.posthogProperties,
     $ai_tokens_source: getTokensSource(params.posthogProperties),
     ...(distinctId ? {} : { $process_person_profile: false }),
+    ...(stopReason ? { $ai_stop_reason: stopReason } : {}),
     ...(tools ? { $ai_tools: tools } : {}),
     ...errorData,
     ...costOverrideData,

packages/ai/src/vercel/middleware.ts

@@ -467,6 +467,15 @@ export const wrapVercelLanguageModel = <T extends LanguageModel>(
 
           adjustAnthropicV3CacheTokens(model, provider, usage)
 
+          // Extract finish reason - V2 returns a string, V3 returns an object with .unified
+          const rawFinishReason = result.finishReason
+          const finishReasonStr =
+            typeof rawFinishReason === 'string'
+              ? rawFinishReason
+              : rawFinishReason && typeof rawFinishReason === 'object' && 'unified' in rawFinishReason
+                ? String(rawFinishReason.unified)
+                : undefined
+
           await sendEventToPosthog({
             client: phClient,
             distinctId: mergedOptions.posthogDistinctId,
@@ -480,6 +489,7 @@ export const wrapVercelLanguageModel = <T extends LanguageModel>(
             params: mergedParams as any,
             httpStatus: 200,
             usage,
+            stopReason: finishReasonStr,
             tools: availableTools,
             captureImmediate: mergedOptions.posthogCaptureImmediate,
           })
@@ -519,6 +529,7 @@ export const wrapVercelLanguageModel = <T extends LanguageModel>(
         let firstTokenTime: number | undefined
         let generatedText = ''
         let reasoningText = ''
+        let stopReason: string | undefined
         let usage: {
           inputTokens?: number
           outputTokens?: number
@@ -610,6 +621,14 @@ export const wrapVercelLanguageModel = <T extends LanguageModel>(
                   cacheReadInputTokens: extractCacheReadTokens(chunkUsage),
                   ...additionalTokenValues,
                 }
+
+                // Extract finish reason - V2 returns a string, V3 returns an object with .unified
+                const rawFinishReason = chunk.finishReason
+                if (typeof rawFinishReason === 'string') {
+                  stopReason = rawFinishReason
+                } else if (rawFinishReason && typeof rawFinishReason === 'object' && 'unified' in rawFinishReason) {
+                  stopReason = String(rawFinishReason.unified)
+                }
               }
               controller.enqueue(chunk)
             },
@@ -676,6 +695,7 @@ export const wrapVercelLanguageModel = <T extends LanguageModel>(
                 params: mergedParams as any,
                 httpStatus: 200,
                 usage: finalUsage,
+                stopReason,
                 tools: availableTools,
                 captureImmediate: mergedOptions.posthogCaptureImmediate,
               })