@posthog/ai 7.2.2 → 7.3.1

package/dist/index.mjs

@@ -2,11 +2,10 @@ import { OpenAI, AzureOpenAI } from 'openai';
 import { Buffer } from 'buffer';
 import * as uuid from 'uuid';
 import { v4 } from 'uuid';
-import { wrapLanguageModel } from 'ai';
 import AnthropicOriginal from '@anthropic-ai/sdk';
 import { GoogleGenAI } from '@google/genai';
 
-var version = "7.2.2";
+var version = "7.3.1";
 
 // Type guards for safer type checking
 const isString = value => {
@@ -859,6 +858,7 @@ let WrappedCompletions$1 = class WrappedCompletions extends Completions {
             try {
               const contentBlocks = [];
               let accumulatedContent = '';
+              let modelFromResponse;
               let usage = {
                 inputTokens: 0,
                 outputTokens: 0,
@@ -867,6 +867,10 @@ let WrappedCompletions$1 = class WrappedCompletions extends Completions {
               // Map to track in-progress tool calls
               const toolCallsInProgress = new Map();
               for await (const chunk of stream1) {
+                // Extract model from chunk (Chat Completions chunks have model field)
+                if (!modelFromResponse && chunk.model) {
+                  modelFromResponse = chunk.model;
+                }
                 const choice = chunk?.choices?.[0];
                 const chunkWebSearchCount = calculateWebSearchCount(chunk);
                 if (chunkWebSearchCount > 0 && chunkWebSearchCount > (usage.webSearchCount ?? 0)) {
@@ -954,7 +958,7 @@ let WrappedCompletions$1 = class WrappedCompletions extends Completions {
               await sendEventToPosthog({
                 client: this.phClient,
                 ...posthogParams,
-                model: openAIParams.model,
+                model: openAIParams.model ?? modelFromResponse,
                 provider: 'openai',
                 input: sanitizeOpenAI(openAIParams.messages),
                 output: formattedOutput,
@@ -1007,7 +1011,7 @@ let WrappedCompletions$1 = class WrappedCompletions extends Completions {
           await sendEventToPosthog({
             client: this.phClient,
             ...posthogParams,
-            model: openAIParams.model,
+            model: openAIParams.model ?? result.model,
             provider: 'openai',
             input: sanitizeOpenAI(openAIParams.messages),
             output: formattedOutput,
@@ -1031,7 +1035,7 @@ let WrappedCompletions$1 = class WrappedCompletions extends Completions {
         await sendEventToPosthog({
           client: this.phClient,
           ...posthogParams,
-          model: String(openAIParams.model ?? ''),
+          model: openAIParams.model,
           provider: 'openai',
           input: sanitizeOpenAI(openAIParams.messages),
           output: [],
@@ -1073,6 +1077,7 @@ let WrappedResponses$1 = class WrappedResponses extends Responses {
           (async () => {
             try {
               let finalContent = [];
+              let modelFromResponse;
               let usage = {
                 inputTokens: 0,
                 outputTokens: 0,
@@ -1080,6 +1085,10 @@ let WrappedResponses$1 = class WrappedResponses extends Responses {
               };
               for await (const chunk of stream1) {
                 if ('response' in chunk && chunk.response) {
+                  // Extract model from response object in chunk (for stored prompts)
+                  if (!modelFromResponse && chunk.response.model) {
+                    modelFromResponse = chunk.response.model;
+                  }
                   const chunkWebSearchCount = calculateWebSearchCount(chunk.response);
                   if (chunkWebSearchCount > 0 && chunkWebSearchCount > (usage.webSearchCount ?? 0)) {
                     usage.webSearchCount = chunkWebSearchCount;
@@ -1103,8 +1112,7 @@ let WrappedResponses$1 = class WrappedResponses extends Responses {
               await sendEventToPosthog({
                 client: this.phClient,
                 ...posthogParams,
-                //@ts-expect-error
-                model: openAIParams.model,
+                model: openAIParams.model ?? modelFromResponse,
                 provider: 'openai',
                 input: formatOpenAIResponsesInput(openAIParams.input, openAIParams.instructions),
                 output: finalContent,
@@ -1126,7 +1134,6 @@ let WrappedResponses$1 = class WrappedResponses extends Responses {
               await sendEventToPosthog({
                 client: this.phClient,
                 ...posthogParams,
-                //@ts-expect-error
                 model: openAIParams.model,
                 provider: 'openai',
                 input: formatOpenAIResponsesInput(openAIParams.input, openAIParams.instructions),
@@ -1159,8 +1166,7 @@ let WrappedResponses$1 = class WrappedResponses extends Responses {
           await sendEventToPosthog({
             client: this.phClient,
             ...posthogParams,
-            //@ts-expect-error
-            model: openAIParams.model,
+            model: openAIParams.model ?? result.model,
             provider: 'openai',
             input: formatOpenAIResponsesInput(openAIParams.input, openAIParams.instructions),
             output: formattedOutput,
@@ -1184,7 +1190,7 @@ let WrappedResponses$1 = class WrappedResponses extends Responses {
         await sendEventToPosthog({
           client: this.phClient,
           ...posthogParams,
-          model: String(openAIParams.model ?? ''),
+          model: openAIParams.model,
           provider: 'openai',
           input: formatOpenAIResponsesInput(openAIParams.input, openAIParams.instructions),
           output: [],
@@ -1221,7 +1227,7 @@ let WrappedResponses$1 = class WrappedResponses extends Responses {
         await sendEventToPosthog({
           client: this.phClient,
           ...posthogParams,
-          model: String(openAIParams.model ?? ''),
+          model: openAIParams.model ?? result.model,
           provider: 'openai',
           input: formatOpenAIResponsesInput(openAIParams.input, openAIParams.instructions),
           output: result.output,
@@ -1242,7 +1248,7 @@ let WrappedResponses$1 = class WrappedResponses extends Responses {
         await sendEventToPosthog({
           client: this.phClient,
           ...posthogParams,
-          model: String(openAIParams.model ?? ''),
+          model: openAIParams.model,
           provider: 'openai',
           input: formatOpenAIResponsesInput(openAIParams.input, openAIParams.instructions),
           output: [],
@@ -1417,7 +1423,7 @@ class WrappedTranscriptions extends Transcriptions {
           await sendEventToPosthog({
             client: this.phClient,
             ...posthogParams,
-            model: String(openAIParams.model ?? ''),
+            model: openAIParams.model,
             provider: 'openai',
             input: openAIParams.prompt,
             output: result.text,
@@ -1437,7 +1443,7 @@ class WrappedTranscriptions extends Transcriptions {
         await sendEventToPosthog({
           client: this.phClient,
           ...posthogParams,
-          model: String(openAIParams.model ?? ''),
+          model: openAIParams.model,
           provider: 'openai',
           input: openAIParams.prompt,
           output: [],
@@ -1499,6 +1505,7 @@ class WrappedCompletions extends AzureOpenAI.Chat.Completions {
             try {
               const contentBlocks = [];
               let accumulatedContent = '';
+              let modelFromResponse;
               let usage = {
                 inputTokens: 0,
                 outputTokens: 0
@@ -1506,6 +1513,10 @@ class WrappedCompletions extends AzureOpenAI.Chat.Completions {
               // Map to track in-progress tool calls
               const toolCallsInProgress = new Map();
               for await (const chunk of stream1) {
+                // Extract model from response if not in params
+                if (!modelFromResponse && chunk.model) {
+                  modelFromResponse = chunk.model;
+                }
                 const choice = chunk?.choices?.[0];
                 // Handle text content
                 const deltaContent = choice?.delta?.content;
@@ -1587,7 +1598,7 @@ class WrappedCompletions extends AzureOpenAI.Chat.Completions {
               await sendEventToPosthog({
                 client: this.phClient,
                 ...posthogParams,
-                model: openAIParams.model,
+                model: openAIParams.model ?? modelFromResponse,
                 provider: 'azure',
                 input: sanitizeOpenAI(openAIParams.messages),
                 output: formattedOutput,
@@ -1631,7 +1642,7 @@ class WrappedCompletions extends AzureOpenAI.Chat.Completions {
           await sendEventToPosthog({
             client: this.phClient,
             ...posthogParams,
-            model: openAIParams.model,
+            model: openAIParams.model ?? result.model,
             provider: 'azure',
             input: openAIParams.messages,
             output: formatResponseOpenAI(result),
@@ -1695,11 +1706,18 @@ class WrappedResponses extends AzureOpenAI.Responses {
           (async () => {
             try {
               let finalContent = [];
+              let modelFromResponse;
               let usage = {
                 inputTokens: 0,
                 outputTokens: 0
               };
               for await (const chunk of stream1) {
+                if ('response' in chunk && chunk.response) {
+                  // Extract model from response if not in params (for stored prompts)
+                  if (!modelFromResponse && chunk.response.model) {
+                    modelFromResponse = chunk.response.model;
+                  }
+                }
                 if (chunk.type === 'response.completed' && 'response' in chunk && chunk.response?.output && chunk.response.output.length > 0) {
                   finalContent = chunk.response.output;
                 }
@@ -1716,10 +1734,9 @@ class WrappedResponses extends AzureOpenAI.Responses {
               await sendEventToPosthog({
                 client: this.phClient,
                 ...posthogParams,
-                //@ts-expect-error
-                model: openAIParams.model,
+                model: openAIParams.model ?? modelFromResponse,
                 provider: 'azure',
-                input: openAIParams.input,
+                input: formatOpenAIResponsesInput(openAIParams.input, openAIParams.instructions),
                 output: finalContent,
                 latency,
                 baseURL: this.baseURL,
@@ -1732,10 +1749,9 @@ class WrappedResponses extends AzureOpenAI.Responses {
               await sendEventToPosthog({
                 client: this.phClient,
                 ...posthogParams,
-                //@ts-expect-error
                 model: openAIParams.model,
                 provider: 'azure',
-                input: openAIParams.input,
+                input: formatOpenAIResponsesInput(openAIParams.input, openAIParams.instructions),
                 output: [],
                 latency: 0,
                 baseURL: this.baseURL,
@@ -1761,10 +1777,9 @@ class WrappedResponses extends AzureOpenAI.Responses {
           await sendEventToPosthog({
             client: this.phClient,
             ...posthogParams,
-            //@ts-expect-error
-            model: openAIParams.model,
+            model: openAIParams.model ?? result.model,
             provider: 'azure',
-            input: openAIParams.input,
+            input: formatOpenAIResponsesInput(openAIParams.input, openAIParams.instructions),
             output: result.output,
             latency,
             baseURL: this.baseURL,
@@ -1784,10 +1799,9 @@ class WrappedResponses extends AzureOpenAI.Responses {
         await sendEventToPosthog({
           client: this.phClient,
           ...posthogParams,
-          //@ts-expect-error
           model: openAIParams.model,
           provider: 'azure',
-          input: openAIParams.input,
+          input: formatOpenAIResponsesInput(openAIParams.input, openAIParams.instructions),
           output: [],
           latency: 0,
           baseURL: this.baseURL,
@@ -1817,9 +1831,9 @@ class WrappedResponses extends AzureOpenAI.Responses {
       await sendEventToPosthog({
         client: this.phClient,
         ...posthogParams,
-        model: String(openAIParams.model ?? ''),
+        model: openAIParams.model ?? result.model,
         provider: 'azure',
-        input: openAIParams.input,
+        input: formatOpenAIResponsesInput(openAIParams.input, openAIParams.instructions),
         output: result.output,
         latency,
         baseURL: this.baseURL,
@@ -1837,9 +1851,9 @@ class WrappedResponses extends AzureOpenAI.Responses {
       await sendEventToPosthog({
         client: this.phClient,
         ...posthogParams,
-        model: String(openAIParams.model ?? ''),
+        model: openAIParams.model,
         provider: 'azure',
-        input: openAIParams.input,
+        input: formatOpenAIResponsesInput(openAIParams.input, openAIParams.instructions),
         output: [],
         latency: 0,
         baseURL: this.baseURL,
@@ -2112,67 +2126,117 @@ const extractProvider = model => {
   const providerName = provider.split('.')[0];
   return providerName;
 };
-const createInstrumentationMiddleware = (phClient, model, options) => {
-  const middleware = {
-    wrapGenerate: async ({
-      doGenerate,
-      params
-    }) => {
+// Extract web search count from provider metadata (works for both V2 and V3)
+const extractWebSearchCount = (providerMetadata, usage) => {
+  // Try Anthropic-specific extraction
+  if (providerMetadata && typeof providerMetadata === 'object' && 'anthropic' in providerMetadata && providerMetadata.anthropic && typeof providerMetadata.anthropic === 'object' && 'server_tool_use' in providerMetadata.anthropic) {
+    const serverToolUse = providerMetadata.anthropic.server_tool_use;
+    if (serverToolUse && typeof serverToolUse === 'object' && 'web_search_requests' in serverToolUse && typeof serverToolUse.web_search_requests === 'number') {
+      return serverToolUse.web_search_requests;
+    }
+  }
+  // Fall back to generic calculation
+  return calculateWebSearchCount({
+    usage,
+    providerMetadata
+  });
+};
+// Extract additional token values from provider metadata
+const extractAdditionalTokenValues = providerMetadata => {
+  if (providerMetadata && typeof providerMetadata === 'object' && 'anthropic' in providerMetadata && providerMetadata.anthropic && typeof providerMetadata.anthropic === 'object' && 'cacheCreationInputTokens' in providerMetadata.anthropic) {
+    return {
+      cacheCreationInputTokens: providerMetadata.anthropic.cacheCreationInputTokens
+    };
+  }
+  return {};
+};
+// Helper to extract numeric token value from V2 (number) or V3 (object with .total) usage formats
+const extractTokenCount = value => {
+  if (typeof value === 'number') {
+    return value;
+  }
+  if (value && typeof value === 'object' && 'total' in value && typeof value.total === 'number') {
+    return value.total;
+  }
+  return undefined;
+};
+// Helper to extract reasoning tokens from V2 (usage.reasoningTokens) or V3 (usage.outputTokens.reasoning)
+const extractReasoningTokens = usage => {
+  // V2 style: top-level reasoningTokens
+  if ('reasoningTokens' in usage) {
+    return usage.reasoningTokens;
+  }
+  // V3 style: nested in outputTokens.reasoning
+  if ('outputTokens' in usage && usage.outputTokens && typeof usage.outputTokens === 'object' && 'reasoning' in usage.outputTokens) {
+    return usage.outputTokens.reasoning;
+  }
+  return undefined;
+};
+// Helper to extract cached input tokens from V2 (usage.cachedInputTokens) or V3 (usage.inputTokens.cacheRead)
+const extractCacheReadTokens = usage => {
+  // V2 style: top-level cachedInputTokens
+  if ('cachedInputTokens' in usage) {
+    return usage.cachedInputTokens;
+  }
+  // V3 style: nested in inputTokens.cacheRead
+  if ('inputTokens' in usage && usage.inputTokens && typeof usage.inputTokens === 'object' && 'cacheRead' in usage.inputTokens) {
+    return usage.inputTokens.cacheRead;
+  }
+  return undefined;
+};
+/**
+ * Wraps a Vercel AI SDK language model (V2 or V3) with PostHog tracing.
+ * Automatically detects the model version and applies appropriate instrumentation.
+ */
+const wrapVercelLanguageModel = (model, phClient, options) => {
+  const traceId = options.posthogTraceId ?? v4();
+  const mergedOptions = {
+    ...options,
+    posthogTraceId: traceId,
+    posthogDistinctId: options.posthogDistinctId,
+    posthogProperties: {
+      ...options.posthogProperties,
+      $ai_framework: 'vercel',
+      $ai_framework_version: model.specificationVersion === 'v3' ? '6' : '5'
+    }
+  };
+  // Create wrapped model that preserves the original type
+  const wrappedModel = {
+    ...model,
+    doGenerate: async params => {
       const startTime = Date.now();
       const mergedParams = {
-        ...options,
-        ...mapVercelParams(params),
-        posthogProperties: {
-          ...options.posthogProperties,
-          $ai_framework: 'vercel'
-        }
+        ...mergedOptions,
+        ...mapVercelParams(params)
       };
       const availableTools = extractAvailableToolCalls('vercel', params);
       try {
-        const result = await doGenerate();
-        const modelId = options.posthogModelOverride ?? (result.response?.modelId ? result.response.modelId : model.modelId);
-        const provider = options.posthogProviderOverride ?? extractProvider(model);
+        const result = await model.doGenerate(params);
+        const modelId = mergedOptions.posthogModelOverride ?? (result.response?.modelId ? result.response.modelId : model.modelId);
+        const provider = mergedOptions.posthogProviderOverride ?? extractProvider(model);
         const baseURL = ''; // cannot currently get baseURL from vercel
         const content = mapVercelOutput(result.content);
         const latency = (Date.now() - startTime) / 1000;
         const providerMetadata = result.providerMetadata;
-        const additionalTokenValues = {
-          ...(providerMetadata?.anthropic ? {
-            cacheCreationInputTokens: providerMetadata.anthropic.cacheCreationInputTokens
-          } : {})
-        };
-        // Calculate web search count based on provider
-        let webSearchCount = 0;
-        if (providerMetadata?.anthropic && typeof providerMetadata.anthropic === 'object' && 'server_tool_use' in providerMetadata.anthropic) {
-          // Anthropic-specific extraction
-          const serverToolUse = providerMetadata.anthropic.server_tool_use;
-          if (serverToolUse && typeof serverToolUse === 'object' && 'web_search_requests' in serverToolUse && typeof serverToolUse.web_search_requests === 'number') {
-            webSearchCount = serverToolUse.web_search_requests;
-          }
-        } else {
-          // For other providers through Vercel, pass available metadata to helper
-          // Note: Vercel abstracts provider responses, so we may not have access to
-          // raw citations/annotations unless Vercel exposes them in usage/metadata
-          webSearchCount = calculateWebSearchCount({
-            usage: result.usage,
-            providerMetadata: providerMetadata
-          });
-        }
+        const additionalTokenValues = extractAdditionalTokenValues(providerMetadata);
+        const webSearchCount = extractWebSearchCount(providerMetadata, result.usage);
+        // V2 usage has simple numbers, V3 has objects with .total - normalize both
+        const usageObj = result.usage;
         const usage = {
-          inputTokens: result.usage.inputTokens,
-          outputTokens: result.usage.outputTokens,
-          reasoningTokens: result.usage.reasoningTokens,
-          cacheReadInputTokens: result.usage.cachedInputTokens,
+          inputTokens: extractTokenCount(result.usage.inputTokens),
+          outputTokens: extractTokenCount(result.usage.outputTokens),
+          reasoningTokens: extractReasoningTokens(usageObj),
+          cacheReadInputTokens: extractCacheReadTokens(usageObj),
           webSearchCount,
           ...additionalTokenValues
         };
         await sendEventToPosthog({
           client: phClient,
-          distinctId: options.posthogDistinctId,
-          traceId: options.posthogTraceId ?? v4(),
+          distinctId: mergedOptions.posthogDistinctId,
+          traceId: mergedOptions.posthogTraceId ?? v4(),
           model: modelId,
           provider: provider,
-          input: options.posthogPrivacyMode ? '' : mapVercelPrompt(params.prompt),
+          input: mergedOptions.posthogPrivacyMode ? '' : mapVercelPrompt(params.prompt),
           output: content,
           latency,
           baseURL,
@@ -2180,18 +2244,18 @@ const createInstrumentationMiddleware = (phClient, model, options) => {
           httpStatus: 200,
           usage,
           tools: availableTools,
-          captureImmediate: options.posthogCaptureImmediate
+          captureImmediate: mergedOptions.posthogCaptureImmediate
         });
         return result;
       } catch (error) {
         const modelId = model.modelId;
         await sendEventToPosthog({
           client: phClient,
-          distinctId: options.posthogDistinctId,
-          traceId: options.posthogTraceId ?? v4(),
+          distinctId: mergedOptions.posthogDistinctId,
+          traceId: mergedOptions.posthogTraceId ?? v4(),
           model: modelId,
           provider: model.provider,
-          input: options.posthogPrivacyMode ? '' : mapVercelPrompt(params.prompt),
+          input: mergedOptions.posthogPrivacyMode ? '' : mapVercelPrompt(params.prompt),
           output: [],
           latency: 0,
           baseURL: '',
@@ -2204,30 +2268,23 @@ const createInstrumentationMiddleware = (phClient, model, options) => {
           isError: true,
           error: truncate(JSON.stringify(error)),
           tools: availableTools,
-          captureImmediate: options.posthogCaptureImmediate
+          captureImmediate: mergedOptions.posthogCaptureImmediate
         });
         throw error;
       }
     },
-    wrapStream: async ({
-      doStream,
-      params
-    }) => {
+    doStream: async params => {
       const startTime = Date.now();
       let generatedText = '';
       let reasoningText = '';
       let usage = {};
       let providerMetadata = undefined;
       const mergedParams = {
-        ...options,
-        ...mapVercelParams(params),
-        posthogProperties: {
-          ...options.posthogProperties,
-          $ai_framework: 'vercel'
-        }
+        ...mergedOptions,
+        ...mapVercelParams(params)
       };
-      const modelId = options.posthogModelOverride ?? model.modelId;
-      const provider = options.posthogProviderOverride ?? extractProvider(model);
+      const modelId = mergedOptions.posthogModelOverride ?? model.modelId;
+      const provider = mergedOptions.posthogProviderOverride ?? extractProvider(model);
       const availableTools = extractAvailableToolCalls('vercel', params);
       const baseURL = ''; // cannot currently get baseURL from vercel
       // Map to track in-progress tool calls
@@ -2236,15 +2293,15 @@ const createInstrumentationMiddleware = (phClient, model, options) => {
         const {
           stream,
           ...rest
-        } = await doStream();
+        } = await model.doStream(params);
         const transformStream = new TransformStream({
           transform(chunk, controller) {
-            // Handle new v5 streaming patterns
+            // Handle streaming patterns - compatible with both V2 and V3
             if (chunk.type === 'text-delta') {
               generatedText += chunk.delta;
             }
             if (chunk.type === 'reasoning-delta') {
-              reasoningText += chunk.delta; // New in v5
+              reasoningText += chunk.delta;
             }
             // Handle tool call chunks
             if (chunk.type === 'tool-input-start') {
@@ -2264,7 +2321,6 @@ const createInstrumentationMiddleware = (phClient, model, options) => {
             }
             if (chunk.type === 'tool-input-end') {
               // Tool call is complete, keep it in the map for final processing
-              // Nothing specific to do here, the tool call is already complete
             }
             if (chunk.type === 'tool-call') {
               // Direct tool call chunk (complete tool call)
@@ -2276,14 +2332,13 @@ const createInstrumentationMiddleware = (phClient, model, options) => {
             }
             if (chunk.type === 'finish') {
               providerMetadata = chunk.providerMetadata;
-              const additionalTokenValues = providerMetadata && typeof providerMetadata === 'object' && 'anthropic' in providerMetadata && providerMetadata.anthropic && typeof providerMetadata.anthropic === 'object' && 'cacheCreationInputTokens' in providerMetadata.anthropic ? {
-                cacheCreationInputTokens: providerMetadata.anthropic.cacheCreationInputTokens
-              } : {};
+              const additionalTokenValues = extractAdditionalTokenValues(providerMetadata);
+              const chunkUsage = chunk.usage || {};
               usage = {
-                inputTokens: chunk.usage?.inputTokens,
-                outputTokens: chunk.usage?.outputTokens,
-                reasoningTokens: chunk.usage?.reasoningTokens,
-                cacheReadInputTokens: chunk.usage?.cachedInputTokens,
+                inputTokens: extractTokenCount(chunk.usage?.inputTokens),
+                outputTokens: extractTokenCount(chunk.usage?.outputTokens),
+                reasoningTokens: extractReasoningTokens(chunkUsage),
+                cacheReadInputTokens: extractCacheReadTokens(chunkUsage),
                 ...additionalTokenValues
               };
             }
@@ -2323,23 +2378,7 @@ const createInstrumentationMiddleware = (phClient, model, options) => {
               role: 'assistant',
               content: content.length === 1 && content[0].type === 'text' ? content[0].text : content
             }] : [];
-            // Calculate web search count based on provider
-            let webSearchCount = 0;
-            if (providerMetadata && typeof providerMetadata === 'object' && 'anthropic' in providerMetadata && providerMetadata.anthropic && typeof providerMetadata.anthropic === 'object' && 'server_tool_use' in providerMetadata.anthropic) {
-              // Anthropic-specific extraction
-              const serverToolUse = providerMetadata.anthropic.server_tool_use;
-              if (serverToolUse && typeof serverToolUse === 'object' && 'web_search_requests' in serverToolUse && typeof serverToolUse.web_search_requests === 'number') {
-                webSearchCount = serverToolUse.web_search_requests;
-              }
-            } else {
-              // For other providers through Vercel, pass available metadata to helper
-              // Note: Vercel abstracts provider responses, so we may not have access to
-              // raw citations/annotations unless Vercel exposes them in usage/metadata
-              webSearchCount = calculateWebSearchCount({
-                usage: usage,
-                providerMetadata: providerMetadata
-              });
-            }
+            const webSearchCount = extractWebSearchCount(providerMetadata, usage);
             // Update usage with web search count
             const finalUsage = {
               ...usage,
@@ -2347,11 +2386,11 @@ const createInstrumentationMiddleware = (phClient, model, options) => {
             };
             await sendEventToPosthog({
               client: phClient,
-              distinctId: options.posthogDistinctId,
-              traceId: options.posthogTraceId ?? v4(),
+              distinctId: mergedOptions.posthogDistinctId,
+              traceId: mergedOptions.posthogTraceId ?? v4(),
               model: modelId,
               provider: provider,
-              input: options.posthogPrivacyMode ? '' : mapVercelPrompt(params.prompt),
+              input: mergedOptions.posthogPrivacyMode ? '' : mapVercelPrompt(params.prompt),
               output: output,
               latency,
               baseURL,
@@ -2359,7 +2398,7 @@ const createInstrumentationMiddleware = (phClient, model, options) => {
               httpStatus: 200,
               usage: finalUsage,
               tools: availableTools,
-              captureImmediate: options.posthogCaptureImmediate
+              captureImmediate: mergedOptions.posthogCaptureImmediate
             });
           }
         });
@@ -2370,11 +2409,11 @@ const createInstrumentationMiddleware = (phClient, model, options) => {
       } catch (error) {
         await sendEventToPosthog({
           client: phClient,
-          distinctId: options.posthogDistinctId,
-          traceId: options.posthogTraceId ?? v4(),
+          distinctId: mergedOptions.posthogDistinctId,
+          traceId: mergedOptions.posthogTraceId ?? v4(),
           model: modelId,
           provider: provider,
-          input: options.posthogPrivacyMode ? '' : mapVercelPrompt(params.prompt),
+          input: mergedOptions.posthogPrivacyMode ? '' : mapVercelPrompt(params.prompt),
           output: [],
           latency: 0,
           baseURL: '',
@@ -2387,25 +2426,12 @@ const createInstrumentationMiddleware = (phClient, model, options) => {
           isError: true,
           error: truncate(JSON.stringify(error)),
           tools: availableTools,
-          captureImmediate: options.posthogCaptureImmediate
+          captureImmediate: mergedOptions.posthogCaptureImmediate
         });
         throw error;
       }
     }
   };
-  return middleware;
-};
-const wrapVercelLanguageModel = (model, phClient, options) => {
-  const traceId = options.posthogTraceId ?? v4();
-  const middleware = createInstrumentationMiddleware(phClient, model, {
-    ...options,
-    posthogTraceId: traceId,
-    posthogDistinctId: options.posthogDistinctId
-  });
-  const wrappedModel = wrapLanguageModel({
-    model,
-    middleware
-  });
   return wrappedModel;
 };