diff --git a/.changeset/add-ai-stop-reason.md b/.changeset/add-ai-stop-reason.md new file mode 100644 index 0000000000..3de625a542 --- /dev/null +++ b/.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 diff --git a/packages/ai/src/anthropic/index.ts b/packages/ai/src/anthropic/index.ts index 9c294e1ae9..47f4eb4e42 100644 --- a/packages/ai/src/anthropic/index.ts +++ b/packages/ai/src/anthropic/index.ts @@ -79,6 +79,7 @@ export class WrappedMessages extends AnthropicOriginal.Messages { const toolsInProgress: Map = 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, }) } diff --git a/packages/ai/src/gemini/index.ts b/packages/ai/src/gemini/index.ts index c016ae7fd2..00e07c782f 100644 --- a/packages/ai/src/gemini/index.ts +++ b/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 { + 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 = [] diff --git a/packages/ai/src/langchain/callbacks.ts b/packages/ai/src/langchain/callbacks.ts index 42710648cc..c9a9653bba 100644 --- a/packages/ai/src/langchain/callbacks.ts +++ b/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 } + 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] { const conversionList: Array<[string, 'input' | 'output']> = [ ['promptTokens', 'input'], diff --git a/packages/ai/src/openai/index.ts b/packages/ai/src/openai/index.ts index f5bbe124de..72d7f1dea8 100644 --- a/packages/ai/src/openai/index.ts +++ b/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 }, diff --git a/packages/ai/src/utils.ts b/packages/ai/src/utils.ts index a5a13e94ea..be06b519d8 100644 --- a/packages/ai/src/utils.ts +++ b/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 => { @@ -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, diff --git a/packages/ai/src/vercel/middleware.ts b/packages/ai/src/vercel/middleware.ts index a629de90fe..98fdd79b40 100644 --- a/packages/ai/src/vercel/middleware.ts +++ b/packages/ai/src/vercel/middleware.ts @@ -467,6 +467,15 @@ export const wrapVercelLanguageModel = ( 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 = ( params: mergedParams as any, httpStatus: 200, usage, + stopReason: finishReasonStr, tools: availableTools, captureImmediate: mergedOptions.posthogCaptureImmediate, }) @@ -519,6 +529,7 @@ export const wrapVercelLanguageModel = ( 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 = ( 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 = ( params: mergedParams as any, httpStatus: 200, usage: finalUsage, + stopReason, tools: availableTools, captureImmediate: mergedOptions.posthogCaptureImmediate, }) diff --git a/packages/ai/tests/gemini.integration.test.ts b/packages/ai/tests/gemini.integration.test.ts new file mode 100644 index 0000000000..06b4ae0c22 --- /dev/null +++ b/packages/ai/tests/gemini.integration.test.ts @@ -0,0 +1,59 @@ +// Integration tests for Gemini generateContent. +// These tests require a real GEMINI_API_KEY and proper ESM transform configuration. +// They are skipped entirely when no API key is present. +// +// To run: GEMINI_API_KEY= jest --testPathPattern=gemini.integration \ +// --transformIgnorePatterns='node_modules/(?!(@google/genai|p-retry|is-network-error)/)' + +const GEMINI_API_KEY = process.env.GEMINI_API_KEY + +if (!GEMINI_API_KEY) { + test.skip('Gemini integration tests require GEMINI_API_KEY', () => {}) +} else { + // Dynamic imports to avoid ESM parse failures when @google/genai + // transitive deps are not configured in transformIgnorePatterns. + // eslint-disable-next-line @typescript-eslint/no-var-requires, @typescript-eslint/no-require-imports + const { PostHog } = require('posthog-node') + + jest.mock('posthog-node', () => ({ + PostHog: jest.fn().mockImplementation(() => ({ + capture: jest.fn(), + captureImmediate: jest.fn(), + privacyMode: false, + })), + })) + + // eslint-disable-next-line @typescript-eslint/no-var-requires, @typescript-eslint/no-require-imports + const PostHogGemini = require('../src/gemini').default + + describe('Gemini Integration Tests', () => { + let mockPostHogClient: any + let client: any + + beforeEach(() => { + mockPostHogClient = new PostHog('test-key') + client = new PostHogGemini({ + apiKey: GEMINI_API_KEY, + posthog: mockPostHogClient, + }) + }) + + test('generateContent captures stop_reason', async () => { + const response = await client.models.generateContent({ + model: 'gemini-2.0-flash', + contents: 'Say hi', + posthogDistinctId: 'test-id', + }) + + expect(response.text).toBeDefined() + + const captureCall = mockPostHogClient.capture.mock.calls.find((call: any[]) => call[0].event === '$ai_generation') + expect(captureCall).toBeDefined() + const props = captureCall![0].properties + expect(props.$ai_provider).toBe('gemini') + expect(props.$ai_stop_reason).toBeDefined() + expect(typeof props.$ai_stop_reason).toBe('string') + expect(props.$ai_input_tokens).toBeGreaterThan(0) + }) + }) +} diff --git a/packages/ai/tests/gemini.test.ts b/packages/ai/tests/gemini.test.ts index d8b38c2903..d253e338fd 100644 --- a/packages/ai/tests/gemini.test.ts +++ b/packages/ai/tests/gemini.test.ts @@ -24,6 +24,7 @@ jest.mock('@google/genai', () => { this.models = { generateContent: jest.fn(), generateContentStream: jest.fn(), + embedContent: jest.fn(), } } } @@ -920,4 +921,141 @@ describe('PostHogGemini - Jest test suite', () => { expect(generateContentCall.config.speechConfig.voiceConfig.prebuiltVoiceConfig.voiceName).toBe('Kore') }) }) + + describe('embedContent', () => { + test('basic embedding', async () => { + const mockEmbedResponse = { + embeddings: [{ values: [0.1, 0.2, 0.3] }], + } + ;(client as any).client.models.embedContent = jest.fn().mockResolvedValue(mockEmbedResponse) + + const response = await client.models.embedContent({ + model: 'gemini-embedding-001', + contents: 'Hello world', + posthogDistinctId: 'test-id', + posthogProperties: { foo: 'bar' }, + }) + + expect(response).toEqual(mockEmbedResponse) + 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('gemini') + expect(properties['$ai_model']).toBe('gemini-embedding-001') + expect(properties['$ai_input']).toBe('Hello world') + expect(properties['$ai_output_choices']).toBeNull() + expect(properties['$ai_http_status']).toBe(200) + expect(properties['foo']).toBe('bar') + expect(typeof properties['$ai_latency']).toBe('number') + expect(properties['$ai_latency']).toBeGreaterThanOrEqual(0) + }) + + test('with token counts from statistics', async () => { + const mockEmbedResponse = { + embeddings: [ + { values: [0.1, 0.2], statistics: { tokenCount: 5 } }, + { values: [0.3, 0.4], statistics: { tokenCount: 3 } }, + { values: [0.5, 0.6], statistics: { tokenCount: 7 } }, + ], + } + ;(client as any).client.models.embedContent = jest.fn().mockResolvedValue(mockEmbedResponse) + + await client.models.embedContent({ + model: 'gemini-embedding-001', + contents: ['Hello', 'World', 'Test'], + posthogDistinctId: 'test-id', + }) + + const [captureArgs] = (mockPostHogClient.capture as jest.Mock).mock.calls + const { properties } = captureArgs[0] + + expect(properties['$ai_input_tokens']).toBe(15) + }) + + test('without token counts defaults to 0', async () => { + const mockEmbedResponse = { + embeddings: [{ values: [0.1, 0.2, 0.3] }], + } + ;(client as any).client.models.embedContent = jest.fn().mockResolvedValue(mockEmbedResponse) + + await client.models.embedContent({ + model: 'gemini-embedding-001', + contents: 'Hello', + posthogDistinctId: 'test-id', + }) + + const [captureArgs] = (mockPostHogClient.capture as jest.Mock).mock.calls + const { properties } = captureArgs[0] + + expect(properties['$ai_input_tokens']).toBe(0) + }) + + test('privacy mode hides input', async () => { + const mockEmbedResponse = { + embeddings: [{ values: [0.1, 0.2] }], + } + ;(client as any).client.models.embedContent = jest.fn().mockResolvedValue(mockEmbedResponse) + + await client.models.embedContent({ + model: 'gemini-embedding-001', + contents: 'Sensitive data', + posthogDistinctId: 'test-id', + posthogPrivacyMode: true, + }) + + const [captureArgs] = (mockPostHogClient.capture as jest.Mock).mock.calls + const { properties } = captureArgs[0] + + expect(properties['$ai_input']).toBeNull() + expect(properties['$ai_output_choices']).toBeNull() + }) + + test('error handling', async () => { + const error = new Error('Embedding API Error') + ;(error as any).status = 429 + ;(client as any).client.models.embedContent = jest.fn().mockRejectedValue(error) + + await expect( + client.models.embedContent({ + model: 'gemini-embedding-001', + contents: 'Hello', + posthogDistinctId: 'test-id', + }) + ).rejects.toThrow('Embedding API Error') + + expect(mockPostHogClient.capture).toHaveBeenCalledTimes(1) + const [captureArgs] = (mockPostHogClient.capture as jest.Mock).mock.calls + const { event, properties } = captureArgs[0] + + expect(event).toBe('$ai_embedding') + expect(properties['$ai_is_error']).toBe(true) + expect(properties['$ai_http_status']).toBe(429) + expect(properties['$ai_input_tokens']).toBe(0) + }) + + test('config passed through to underlying call', async () => { + const mockEmbedResponse = { + embeddings: [{ values: [0.1, 0.2] }], + } + ;(client as any).client.models.embedContent = jest.fn().mockResolvedValue(mockEmbedResponse) + + await client.models.embedContent({ + model: 'gemini-embedding-001', + contents: 'Hello', + config: { outputDimensionality: 256 }, + posthogDistinctId: 'test-id', + } as any) + + const embedCall = ((client as any).client.models.embedContent as jest.Mock).mock.calls[0][0] + expect(embedCall.config).toEqual({ outputDimensionality: 256 }) + expect(embedCall.model).toBe('gemini-embedding-001') + expect(embedCall.contents).toBe('Hello') + // Posthog params should not be passed through + expect(embedCall.posthogDistinctId).toBeUndefined() + }) + }) }) diff --git a/packages/ai/tests/stop-reason.test.ts b/packages/ai/tests/stop-reason.test.ts new file mode 100644 index 0000000000..feed23678a --- /dev/null +++ b/packages/ai/tests/stop-reason.test.ts @@ -0,0 +1,943 @@ +import { PostHog } from 'posthog-node' +import PostHogOpenAI from '../src/openai' +import PostHogAnthropic from '../src/anthropic' +import PostHogGemini from '../src/gemini' +import { LangChainCallbackHandler } from '../src/langchain/callbacks' +import { withTracing } from '../src/index' +import { flushPromises } from './test-utils' +import openaiModule from 'openai' +import AnthropicOriginal from '@anthropic-ai/sdk' + +import type { ChatCompletion, ChatCompletionChunk } from 'openai/resources/chat/completions' +import type { + LanguageModelV2, + LanguageModelV2StreamPart, + LanguageModelV3, + LanguageModelV3StreamPart, +} from '@ai-sdk/provider' + +// --- Mocks --- + +jest.mock('posthog-node', () => ({ + PostHog: jest.fn().mockImplementation(() => ({ + capture: jest.fn(), + captureImmediate: jest.fn(), + privacy_mode: false, + })), +})) + +jest.mock('openai', () => { + class MockCompletions { + create(..._args: any[]): any { + return undefined + } + } + class MockChat { + constructor() {} + static Completions = MockCompletions + } + class MockResponses { + constructor() {} + create() { + return Promise.resolve({}) + } + parse() { + return Promise.resolve({}) + } + } + class MockEmbeddings { + constructor() {} + create() { + return Promise.resolve({}) + } + } + class MockTranscriptions { + constructor() {} + create() { + return Promise.resolve({}) + } + } + class MockAudio { + constructor() {} + static Transcriptions = MockTranscriptions + } + class MockOpenAI { + chat: any + embeddings: any + responses: any + audio: any + constructor() { + this.chat = { completions: { create: jest.fn() } } + this.embeddings = { create: jest.fn() } + this.responses = { create: jest.fn() } + this.audio = { transcriptions: { create: jest.fn() } } + } + static Chat = MockChat + static Responses = MockResponses + static Embeddings = MockEmbeddings + static Audio = MockAudio + } + return { + __esModule: true, + default: MockOpenAI, + OpenAI: MockOpenAI, + AzureOpenAI: MockOpenAI, + Chat: MockChat, + Responses: MockResponses, + Embeddings: MockEmbeddings, + Audio: MockAudio, + } +}) + +jest.mock('@anthropic-ai/sdk', () => { + class MockMessages { + create(..._args: any[]): any { + return undefined + } + } + class MockAnthropic { + messages: any + constructor() { + this.messages = new MockMessages() + } + static Messages = MockMessages + } + return { __esModule: true, default: MockAnthropic } +}) + +jest.mock('@google/genai', () => { + class MockGoogleGenAI { + models: any + constructor() { + this.models = { + generateContent: jest.fn(), + generateContentStream: jest.fn(), + } + } + } + return { GoogleGenAI: MockGoogleGenAI } +}) + +// --- Helpers --- + +interface MockAsyncIterator { + [Symbol.asyncIterator](): AsyncIterator +} + +const createMockAsyncIterator = (chunks: T[]): MockAsyncIterator => { + let index = 0 + return { + async *[Symbol.asyncIterator]() { + while (index < chunks.length) { + yield chunks[index++] + } + }, + } +} + +const getCapturedProperties = (client: PostHog): Record => { + const captureMock = client.capture as jest.Mock + expect(captureMock).toHaveBeenCalledTimes(1) + return captureMock.mock.calls[0][0].properties +} + +// --- Tests --- + +describe('$ai_stop_reason extraction', () => { + let mockPostHogClient: PostHog + + beforeEach(() => { + jest.clearAllMocks() + mockPostHogClient = new (PostHog as any)() + }) + + describe('OpenAI Chat Completions', () => { + let client: PostHogOpenAI + + beforeEach(() => { + client = new PostHogOpenAI({ + apiKey: 'test-key', + posthog: mockPostHogClient as any, + }) + }) + + test('non-streaming: extracts finish_reason', async () => { + const mockResponse: ChatCompletion = { + id: 'test-id', + model: 'gpt-4', + object: 'chat.completion', + created: Date.now() / 1000, + choices: [ + { + index: 0, + finish_reason: 'stop', + message: { role: 'assistant', content: 'Hello!', refusal: null }, + logprobs: null, + }, + ], + usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }, + } + + const ChatMock: any = openaiModule.Chat + ;(ChatMock.Completions as any).prototype.create = jest.fn().mockResolvedValue(mockResponse) + + await client.chat.completions.create({ + model: 'gpt-4', + messages: [{ role: 'user', content: 'Hi' }], + posthogDistinctId: 'test-user', + }) + + const properties = getCapturedProperties(mockPostHogClient) + expect(properties['$ai_stop_reason']).toBe('stop') + }) + + test('non-streaming: extracts length finish_reason', async () => { + const mockResponse: ChatCompletion = { + id: 'test-id', + model: 'gpt-4', + object: 'chat.completion', + created: Date.now() / 1000, + choices: [ + { + index: 0, + finish_reason: 'length', + message: { role: 'assistant', content: 'Truncated...', refusal: null }, + logprobs: null, + }, + ], + usage: { prompt_tokens: 10, completion_tokens: 100, total_tokens: 110 }, + } + + const ChatMock: any = openaiModule.Chat + ;(ChatMock.Completions as any).prototype.create = jest.fn().mockResolvedValue(mockResponse) + + await client.chat.completions.create({ + model: 'gpt-4', + messages: [{ role: 'user', content: 'Write a long essay' }], + posthogDistinctId: 'test-user', + }) + + const properties = getCapturedProperties(mockPostHogClient) + expect(properties['$ai_stop_reason']).toBe('length') + }) + + test('streaming: extracts finish_reason from final chunk', async () => { + const chunks: ChatCompletionChunk[] = [ + { + id: 'test', + model: 'gpt-4', + object: 'chat.completion.chunk', + created: Date.now() / 1000, + choices: [{ index: 0, delta: { content: 'Hello!' }, finish_reason: null, logprobs: null }], + }, + { + id: 'test', + model: 'gpt-4', + object: 'chat.completion.chunk', + created: Date.now() / 1000, + choices: [{ index: 0, delta: {}, finish_reason: 'stop', logprobs: null }], + usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }, + }, + ] + + const ChatMock: any = openaiModule.Chat + ;(ChatMock.Completions as any).prototype.create = jest.fn().mockImplementation(() => { + const mockStream = { + tee: jest.fn().mockReturnValue([createMockAsyncIterator(chunks), createMockAsyncIterator(chunks)]), + } + return Promise.resolve(mockStream) + }) + + const stream = await client.chat.completions.create({ + model: 'gpt-4', + messages: [{ role: 'user', content: 'Hi' }], + stream: true, + posthogDistinctId: 'test-user', + }) + + // Consume the stream + for await (const _chunk of stream as any) { + // consume + } + + await flushPromises() + + const properties = getCapturedProperties(mockPostHogClient) + expect(properties['$ai_stop_reason']).toBe('stop') + }) + }) + + describe('OpenAI Responses API', () => { + let client: PostHogOpenAI + + beforeEach(() => { + client = new PostHogOpenAI({ + apiKey: 'test-key', + posthog: mockPostHogClient as any, + }) + }) + + test('non-streaming: extracts status as stop reason', async () => { + const mockResponse = { + id: 'resp-test', + model: 'gpt-4o', + object: 'response', + status: 'completed', + output: [ + { + id: 'msg-001', + type: 'message', + role: 'assistant', + content: [{ type: 'output_text', text: 'Hello!' }], + }, + ], + usage: { + input_tokens: 10, + output_tokens: 5, + input_tokens_details: { cached_tokens: 0 }, + output_tokens_details: { reasoning_tokens: 0 }, + total_tokens: 15, + }, + } + + const ResponsesMock: any = openaiModule.Responses + ResponsesMock.prototype.create = jest.fn().mockResolvedValue(mockResponse) + + await client.responses.create({ + model: 'gpt-4o', + input: 'Hi', + posthogDistinctId: 'test-user', + }) + + const properties = getCapturedProperties(mockPostHogClient) + expect(properties['$ai_stop_reason']).toBe('completed') + }) + + test('streaming: extracts status from response.completed event', async () => { + const chunks = [ + { + type: 'response.output_item.added', + response: { model: 'gpt-4o' }, + }, + { + type: 'response.output_text.delta', + delta: 'Hello!', + }, + { + type: 'response.completed', + response: { + model: 'gpt-4o', + status: 'completed', + output: [ + { + id: 'msg-001', + type: 'message', + role: 'assistant', + content: [{ type: 'output_text', text: 'Hello!' }], + }, + ], + usage: { + input_tokens: 10, + output_tokens: 5, + input_tokens_details: { cached_tokens: 0 }, + output_tokens_details: { reasoning_tokens: 0 }, + }, + }, + }, + ] + + const ResponsesMock: any = openaiModule.Responses + ResponsesMock.prototype.create = jest.fn().mockImplementation(() => { + const mockStream = { + tee: jest.fn().mockReturnValue([createMockAsyncIterator(chunks), createMockAsyncIterator(chunks)]), + } + return Promise.resolve(mockStream) + }) + + const stream = await client.responses.create({ + model: 'gpt-4o', + input: 'Hi', + stream: true, + posthogDistinctId: 'test-user', + }) + + for await (const _chunk of stream as any) { + // consume + } + + await flushPromises() + + const properties = getCapturedProperties(mockPostHogClient) + expect(properties['$ai_stop_reason']).toBe('completed') + }) + }) + + describe('Anthropic', () => { + let client: PostHogAnthropic + + beforeEach(() => { + client = new PostHogAnthropic({ + apiKey: 'test-key', + posthog: mockPostHogClient as any, + }) + }) + + test('non-streaming: extracts stop_reason', async () => { + const mockResponse = { + id: 'msg_test', + type: 'message', + role: 'assistant', + model: 'claude-3-opus-20240229', + content: [{ type: 'text', text: 'Hello!' }], + stop_reason: 'end_turn', + stop_sequence: null, + usage: { input_tokens: 10, output_tokens: 5 }, + } + + const MessagesMock = AnthropicOriginal.Messages as jest.MockedClass + ;(MessagesMock.prototype.create as jest.Mock) = jest.fn().mockResolvedValue(mockResponse) + + await client.messages.create({ + model: 'claude-3-opus-20240229', + messages: [{ role: 'user', content: 'Hi' }], + max_tokens: 100, + posthogDistinctId: 'test-user', + }) + + const properties = getCapturedProperties(mockPostHogClient) + expect(properties['$ai_stop_reason']).toBe('end_turn') + }) + + test('non-streaming: extracts max_tokens stop_reason', async () => { + const mockResponse = { + id: 'msg_test', + type: 'message', + role: 'assistant', + model: 'claude-3-opus-20240229', + content: [{ type: 'text', text: 'Truncated...' }], + stop_reason: 'max_tokens', + stop_sequence: null, + usage: { input_tokens: 10, output_tokens: 100 }, + } + + const MessagesMock = AnthropicOriginal.Messages as jest.MockedClass + ;(MessagesMock.prototype.create as jest.Mock) = jest.fn().mockResolvedValue(mockResponse) + + await client.messages.create({ + model: 'claude-3-opus-20240229', + messages: [{ role: 'user', content: 'Write a long essay' }], + max_tokens: 100, + posthogDistinctId: 'test-user', + }) + + const properties = getCapturedProperties(mockPostHogClient) + expect(properties['$ai_stop_reason']).toBe('max_tokens') + }) + + test('streaming: extracts stop_reason from message_delta', async () => { + const chunks = [ + { + type: 'message_start', + message: { + id: 'msg_test', + type: 'message', + role: 'assistant', + model: 'claude-3-opus-20240229', + usage: { input_tokens: 10, output_tokens: 0 }, + }, + }, + { + type: 'content_block_start', + index: 0, + content_block: { type: 'text', text: '' }, + }, + { + type: 'content_block_delta', + index: 0, + delta: { type: 'text_delta', text: 'Hello!' }, + }, + { type: 'content_block_stop', index: 0 }, + { + type: 'message_delta', + delta: { type: 'stop_reason', stop_reason: 'end_turn' }, + usage: { output_tokens: 5 }, + }, + { type: 'message_stop' }, + ] + + const MessagesMock = AnthropicOriginal.Messages as jest.MockedClass + ;(MessagesMock.prototype.create as jest.Mock) = jest.fn().mockImplementation(() => { + const mockStream = { + tee: jest.fn().mockReturnValue([createMockAsyncIterator(chunks), createMockAsyncIterator(chunks)]), + } + return Promise.resolve(mockStream) + }) + + const stream = await client.messages.create({ + model: 'claude-3-opus-20240229', + messages: [{ role: 'user', content: 'Hi' }], + max_tokens: 100, + stream: true, + posthogDistinctId: 'test-user', + }) + + for await (const _chunk of stream as any) { + // consume + } + + await new Promise(process.nextTick) + + const properties = getCapturedProperties(mockPostHogClient) + expect(properties['$ai_stop_reason']).toBe('end_turn') + }) + }) + + describe('Gemini', () => { + let client: PostHogGemini + + beforeEach(() => { + client = new PostHogGemini({ + apiKey: 'test-key', + posthog: mockPostHogClient as any, + }) + }) + + test('non-streaming: extracts finishReason', async () => { + const mockResponse = { + text: 'Hello from Gemini!', + candidates: [ + { + content: { parts: [{ text: 'Hello from Gemini!' }] }, + finishReason: 'STOP', + }, + ], + usageMetadata: { + promptTokenCount: 10, + candidatesTokenCount: 5, + totalTokenCount: 15, + }, + } + + ;(client as any).client.models.generateContent = jest.fn().mockResolvedValue(mockResponse) + + await client.models.generateContent({ + model: 'gemini-2.0-flash-001', + contents: 'Hello', + posthogDistinctId: 'test-user', + }) + + const properties = getCapturedProperties(mockPostHogClient) + expect(properties['$ai_stop_reason']).toBe('STOP') + }) + + test('non-streaming: extracts MAX_TOKENS finishReason', async () => { + const mockResponse = { + text: 'Truncated...', + candidates: [ + { + content: { parts: [{ text: 'Truncated...' }] }, + finishReason: 'MAX_TOKENS', + }, + ], + usageMetadata: { + promptTokenCount: 10, + candidatesTokenCount: 100, + totalTokenCount: 110, + }, + } + + ;(client as any).client.models.generateContent = jest.fn().mockResolvedValue(mockResponse) + + await client.models.generateContent({ + model: 'gemini-2.0-flash-001', + contents: 'Write a long essay', + posthogDistinctId: 'test-user', + }) + + const properties = getCapturedProperties(mockPostHogClient) + expect(properties['$ai_stop_reason']).toBe('MAX_TOKENS') + }) + + test('streaming: extracts finishReason from final chunk', async () => { + const streamChunks = [ + { + text: 'Hello ', + candidates: [{ content: { parts: [{ text: 'Hello ' }] } }], + usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 2 }, + }, + { + text: 'world!', + candidates: [ + { + content: { parts: [{ text: 'world!' }] }, + finishReason: 'STOP', + }, + ], + usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 5 }, + }, + ] + + ;(client as any).client.models.generateContentStream = jest.fn().mockImplementation(() => { + return (async function* () { + for (const chunk of streamChunks) { + yield chunk + } + })() + }) + + const stream = client.models.generateContentStream({ + model: 'gemini-2.0-flash-001', + contents: 'Hello', + posthogDistinctId: 'test-user', + }) + + for await (const _chunk of stream) { + // consume + } + + const properties = getCapturedProperties(mockPostHogClient) + expect(properties['$ai_stop_reason']).toBe('STOP') + }) + }) + + describe('Vercel AI SDK', () => { + test('V2 doGenerate: extracts finishReason string', async () => { + const baseModel: LanguageModelV2 = { + specificationVersion: 'v2' as const, + provider: 'openai', + modelId: 'gpt-4', + supportedUrls: {}, + doGenerate: jest.fn().mockResolvedValue({ + text: 'Hello!', + usage: { inputTokens: 10, outputTokens: 5 }, + content: [{ type: 'text', text: 'Hello!' }], + response: { modelId: 'gpt-4' }, + providerMetadata: {}, + finishReason: 'stop', + warnings: [], + }), + doStream: jest.fn(), + } + + const model = withTracing(baseModel, mockPostHogClient, { + posthogDistinctId: 'test-user', + }) + + await model.doGenerate({ + prompt: [{ role: 'user', content: [{ type: 'text', text: 'Hi' }] }], + } as any) + + const properties = getCapturedProperties(mockPostHogClient) + expect(properties['$ai_stop_reason']).toBe('stop') + }) + + test('V3 doGenerate: extracts finishReason from object', async () => { + const baseModel: LanguageModelV3 = { + specificationVersion: 'v3' as const, + provider: 'openai', + modelId: 'gpt-4', + supportedUrls: {}, + doGenerate: jest.fn().mockResolvedValue({ + text: 'Hello!', + usage: { + inputTokens: { total: 10, noCache: undefined, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 5, text: undefined, reasoning: undefined }, + }, + content: [{ type: 'text', text: 'Hello!' }], + response: { modelId: 'gpt-4' }, + providerMetadata: {}, + finishReason: { unified: 'stop', raw: undefined }, + warnings: [], + }), + doStream: jest.fn(), + } + + const model = withTracing(baseModel, mockPostHogClient, { + posthogDistinctId: 'test-user', + }) + + await model.doGenerate({ + prompt: [{ role: 'user', content: [{ type: 'text', text: 'Hi' }] }], + } as any) + + const properties = getCapturedProperties(mockPostHogClient) + expect(properties['$ai_stop_reason']).toBe('stop') + }) + + test('V2 doStream: extracts finishReason from finish part', async () => { + const streamParts: LanguageModelV2StreamPart[] = [ + { type: 'text-delta', id: 'text-1', delta: 'Hello!' }, + { + type: 'finish', + usage: { inputTokens: 10, outputTokens: 5, totalTokens: 15 }, + finishReason: 'stop' as const, + }, + ] + + const baseModel: LanguageModelV2 = { + specificationVersion: 'v2' as const, + provider: 'openai', + modelId: 'gpt-4', + supportedUrls: {}, + doGenerate: jest.fn(), + doStream: jest.fn().mockImplementation(async () => { + const stream = new ReadableStream({ + async start(controller) { + for (const part of streamParts) { + controller.enqueue(part) + } + controller.close() + }, + }) + return { stream, response: { modelId: 'gpt-4' } } + }), + } + + const model = withTracing(baseModel, mockPostHogClient, { + posthogDistinctId: 'test-user', + posthogTraceId: 'test-trace', + }) + + const result = await model.doStream({ + prompt: [{ role: 'user' as const, content: [{ type: 'text' as const, text: 'Hi' }] }], + } as any) + + const reader = result.stream.getReader() + while (!(await reader.read()).done) { + // consume + } + + await flushPromises() + + const properties = getCapturedProperties(mockPostHogClient) + expect(properties['$ai_stop_reason']).toBe('stop') + }) + + test('V3 doStream: extracts finishReason from finish part', async () => { + const streamParts: LanguageModelV3StreamPart[] = [ + { type: 'text-delta', id: 'text-1', delta: 'Hello!' }, + { + type: 'finish', + usage: { + inputTokens: { total: 10, noCache: undefined, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 5, text: undefined, reasoning: undefined }, + }, + finishReason: { unified: 'stop' as const, raw: undefined }, + }, + ] + + const baseModel: LanguageModelV3 = { + specificationVersion: 'v3' as const, + provider: 'openai', + modelId: 'gpt-4', + supportedUrls: {}, + doGenerate: jest.fn(), + doStream: jest.fn().mockImplementation(async () => { + const stream = new ReadableStream({ + async start(controller) { + for (const part of streamParts) { + controller.enqueue(part) + } + controller.close() + }, + }) + return { stream, response: { modelId: 'gpt-4' } } + }), + } + + const model = withTracing(baseModel, mockPostHogClient, { + posthogDistinctId: 'test-user', + posthogTraceId: 'test-trace', + }) + + const result = await model.doStream({ + prompt: [{ role: 'user' as const, content: [{ type: 'text' as const, text: 'Hi' }] }], + } as any) + + const reader = result.stream.getReader() + while (!(await reader.read()).done) { + // consume + } + + await flushPromises() + + const properties = getCapturedProperties(mockPostHogClient) + expect(properties['$ai_stop_reason']).toBe('stop') + }) + }) + + describe('LangChain', () => { + test('extracts finish_reason from generationInfo (OpenAI format)', () => { + const handler = new LangChainCallbackHandler({ + client: mockPostHogClient, + }) + + const serialized = { + lc: 1, + type: 'constructor' as const, + id: ['langchain', 'llms', 'openai', 'OpenAI'], + kwargs: {}, + } + + handler.handleLLMStart( + serialized, + ['Test prompt'], + 'run-1', + 'parent-1', + { + invocation_params: { temperature: 0.7 }, + }, + undefined, + { ls_model_name: 'gpt-4', ls_provider: 'openai' } + ) + + handler.handleLLMEnd( + { + generations: [ + [ + { + text: 'Hello!', + + generationInfo: { finish_reason: 'stop' }, + }, + ], + ], + llmOutput: { + tokenUsage: { promptTokens: 10, completionTokens: 5 }, + }, + }, + 'run-1' + ) + + const properties = getCapturedProperties(mockPostHogClient) + expect(properties['$ai_stop_reason']).toBe('stop') + }) + + test('extracts stop_reason from generationInfo (Anthropic format)', () => { + const handler = new LangChainCallbackHandler({ + client: mockPostHogClient, + }) + + const serialized = { + lc: 1, + type: 'constructor' as const, + id: ['langchain', 'llms', 'anthropic'], + kwargs: {}, + } + + handler.handleLLMStart( + serialized, + ['Test prompt'], + 'run-2', + 'parent-2', + { + invocation_params: {}, + }, + undefined, + { ls_model_name: 'claude-3', ls_provider: 'anthropic' } + ) + + handler.handleLLMEnd( + { + generations: [ + [ + { + text: 'Hello!', + + generationInfo: { stop_reason: 'end_turn' }, + }, + ], + ], + llmOutput: { + tokenUsage: { promptTokens: 10, completionTokens: 5 }, + }, + }, + 'run-2' + ) + + const properties = getCapturedProperties(mockPostHogClient) + expect(properties['$ai_stop_reason']).toBe('end_turn') + }) + + test('does not include $ai_stop_reason when not available', () => { + const handler = new LangChainCallbackHandler({ + client: mockPostHogClient, + }) + + const serialized = { + lc: 1, + type: 'constructor' as const, + id: ['langchain', 'llms', 'openai', 'OpenAI'], + kwargs: {}, + } + + handler.handleLLMStart( + serialized, + ['Test prompt'], + 'run-3', + 'parent-3', + { + invocation_params: {}, + }, + undefined, + { ls_model_name: 'gpt-4', ls_provider: 'openai' } + ) + + handler.handleLLMEnd( + { + generations: [ + [ + { + text: 'Hello!', + }, + ], + ], + llmOutput: { + tokenUsage: { promptTokens: 10, completionTokens: 5 }, + }, + }, + 'run-3' + ) + + const properties = getCapturedProperties(mockPostHogClient) + expect(properties['$ai_stop_reason']).toBeUndefined() + }) + }) + + describe('stop reason is omitted when absent', () => { + test('OpenAI non-streaming: no stop reason when finish_reason is null', async () => { + const mockResponse: ChatCompletion = { + id: 'test-id', + model: 'gpt-4', + object: 'chat.completion', + created: Date.now() / 1000, + choices: [ + { + index: 0, + finish_reason: null as any, + message: { role: 'assistant', content: 'Hello!', refusal: null }, + logprobs: null, + }, + ], + usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }, + } + + const client = new PostHogOpenAI({ + apiKey: 'test-key', + posthog: mockPostHogClient as any, + }) + + const ChatMock: any = openaiModule.Chat + ;(ChatMock.Completions as any).prototype.create = jest.fn().mockResolvedValue(mockResponse) + + await client.chat.completions.create({ + model: 'gpt-4', + messages: [{ role: 'user', content: 'Hi' }], + posthogDistinctId: 'test-user', + }) + + const properties = getCapturedProperties(mockPostHogClient) + expect(properties['$ai_stop_reason']).toBeUndefined() + }) + }) +})