@posthog/ai 7.2.2 → 7.3.1

package/dist/vercel/index.mjs

@@ -1,8 +1,7 @@
-import { wrapLanguageModel } from 'ai';
 import { v4 } from 'uuid';
 import { Buffer } from 'buffer';
 
-var version = "7.2.2";
+var version = "7.3.1";
 
 // Type guards for safer type checking
 
@@ -582,68 +581,126 @@ 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
-          } : {})
-        };
+        const additionalTokenValues = extractAdditionalTokenValues(providerMetadata);
+        const webSearchCount = extractWebSearchCount(providerMetadata, result.usage);
 
-        // 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
-          });
-        }
+        // 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,
@@ -651,18 +708,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: '',
@@ -675,30 +732,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
 
@@ -708,15 +758,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
@@ -737,7 +787,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)
@@ -749,14 +798,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
               };
             }
@@ -798,24 +846,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 = {
@@ -824,11 +855,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,
@@ -836,7 +867,7 @@ const createInstrumentationMiddleware = (phClient, model, options) => {
               httpStatus: 200,
               usage: finalUsage,
               tools: availableTools,
-              captureImmediate: options.posthogCaptureImmediate
+              captureImmediate: mergedOptions.posthogCaptureImmediate
             });
           }
         });
@@ -847,11 +878,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: '',
@@ -864,25 +895,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;
 };