@posthog/ai 6.1.1 → 6.2.0

package/dist/index.mjs

@@ -6,6 +6,8 @@ import { wrapLanguageModel } from 'ai';
 import AnthropicOriginal from '@anthropic-ai/sdk';
 import { GoogleGenAI } from '@google/genai';
 
+var version = "6.2.0";
+
 // limit large outputs by truncating to 200kb (approx 200k bytes)
 const MAX_OUTPUT_SIZE = 200000;
 const STRING_FORMAT = 'utf8';
@@ -238,9 +240,8 @@ const extractAvailableToolCalls = (provider, params) => {
     }
     return null;
   } else if (provider === 'vercel') {
-    // Vercel AI SDK stores tools in params.mode.tools when mode type is 'regular'
-    if (params.mode?.type === 'regular' && params.mode.tools) {
-      return params.mode.tools;
+    if (params.tools) {
+      return params.tools;
     }
     return null;
   }
@@ -314,6 +315,8 @@ const sendEventToPosthog = async ({
     } : {})
   };
   const properties = {
+    $ai_lib: 'posthog-ai',
+    $ai_lib_version: version,
     $ai_provider: params.posthogProviderOverride ?? provider,
     $ai_model: params.posthogModelOverride ?? model,
     $ai_model_parameters: getModelParams(params),
@@ -595,14 +598,52 @@ let WrappedCompletions$1 = class WrappedCompletions extends Completions {
           const [stream1, stream2] = value.tee();
           (async () => {
             try {
+              const contentBlocks = [];
               let accumulatedContent = '';
               let usage = {
                 inputTokens: 0,
                 outputTokens: 0
               };
+              // Map to track in-progress tool calls
+              const toolCallsInProgress = new Map();
               for await (const chunk of stream1) {
-                const delta = chunk?.choices?.[0]?.delta?.content ?? '';
-                accumulatedContent += delta;
+                const choice = chunk?.choices?.[0];
+                // Handle text content
+                const deltaContent = choice?.delta?.content;
+                if (deltaContent) {
+                  accumulatedContent += deltaContent;
+                }
+                // Handle tool calls
+                const deltaToolCalls = choice?.delta?.tool_calls;
+                if (deltaToolCalls && Array.isArray(deltaToolCalls)) {
+                  for (const toolCall of deltaToolCalls) {
+                    const index = toolCall.index;
+                    if (index !== undefined) {
+                      if (!toolCallsInProgress.has(index)) {
+                        // New tool call
+                        toolCallsInProgress.set(index, {
+                          id: toolCall.id || '',
+                          name: toolCall.function?.name || '',
+                          arguments: ''
+                        });
+                      }
+                      const inProgressCall = toolCallsInProgress.get(index);
+                      if (inProgressCall) {
+                        // Update tool call data
+                        if (toolCall.id) {
+                          inProgressCall.id = toolCall.id;
+                        }
+                        if (toolCall.function?.name) {
+                          inProgressCall.name = toolCall.function.name;
+                        }
+                        if (toolCall.function?.arguments) {
+                          inProgressCall.arguments += toolCall.function.arguments;
+                        }
+                      }
+                    }
+                  }
+                }
+                // Handle usage information
                 if (chunk.usage) {
                   usage = {
                     inputTokens: chunk.usage.prompt_tokens ?? 0,
@@ -612,6 +653,37 @@ let WrappedCompletions$1 = class WrappedCompletions extends Completions {
                   };
                 }
               }
+              // Build final content blocks
+              if (accumulatedContent) {
+                contentBlocks.push({
+                  type: 'text',
+                  text: accumulatedContent
+                });
+              }
+              // Add completed tool calls to content blocks
+              for (const toolCall of toolCallsInProgress.values()) {
+                if (toolCall.name) {
+                  contentBlocks.push({
+                    type: 'function',
+                    id: toolCall.id,
+                    function: {
+                      name: toolCall.name,
+                      arguments: toolCall.arguments
+                    }
+                  });
+                }
+              }
+              // Format output to match non-streaming version
+              const formattedOutput = contentBlocks.length > 0 ? [{
+                role: 'assistant',
+                content: contentBlocks
+              }] : [{
+                role: 'assistant',
+                content: [{
+                  type: 'text',
+                  text: ''
+                }]
+              }];
               const latency = (Date.now() - startTime) / 1000;
               const availableTools = extractAvailableToolCalls('openai', openAIParams);
               await sendEventToPosthog({
@@ -621,10 +693,7 @@ let WrappedCompletions$1 = class WrappedCompletions extends Completions {
                 model: openAIParams.model,
                 provider: 'openai',
                 input: sanitizeOpenAI(openAIParams.messages),
-                output: [{
-                  content: accumulatedContent,
-                  role: 'assistant'
-                }],
+                output: formattedOutput,
                 latency,
                 baseURL: this.baseURL ?? '',
                 params: body,
@@ -634,6 +703,7 @@ let WrappedCompletions$1 = class WrappedCompletions extends Completions {
                 captureImmediate: posthogCaptureImmediate
               });
             } catch (error) {
+              const httpStatus = error && typeof error === 'object' && 'status' in error ? error.status ?? 500 : 500;
               await sendEventToPosthog({
                 client: this.phClient,
                 distinctId: posthogDistinctId,
@@ -645,7 +715,7 @@ let WrappedCompletions$1 = class WrappedCompletions extends Completions {
                 latency: 0,
                 baseURL: this.baseURL ?? '',
                 params: body,
-                httpStatus: error?.status ? error.status : 500,
+                httpStatus,
                 usage: {
                   inputTokens: 0,
                   outputTokens: 0
@@ -690,6 +760,7 @@ let WrappedCompletions$1 = class WrappedCompletions extends Completions {
         }
         return result;
       }, async error => {
+        const httpStatus = error && typeof error === 'object' && 'status' in error ? error.status ?? 500 : 500;
         await sendEventToPosthog({
           client: this.phClient,
           distinctId: posthogDistinctId,
@@ -701,7 +772,7 @@ let WrappedCompletions$1 = class WrappedCompletions extends Completions {
           latency: 0,
           baseURL: this.baseURL ?? '',
           params: body,
-          httpStatus: error?.status ? error.status : 500,
+          httpStatus,
           usage: {
             inputTokens: 0,
             outputTokens: 0
@@ -780,6 +851,7 @@ let WrappedResponses$1 = class WrappedResponses extends Responses {
                 captureImmediate: posthogCaptureImmediate
               });
             } catch (error) {
+              const httpStatus = error && typeof error === 'object' && 'status' in error ? error.status ?? 500 : 500;
               await sendEventToPosthog({
                 client: this.phClient,
                 distinctId: posthogDistinctId,
@@ -792,7 +864,7 @@ let WrappedResponses$1 = class WrappedResponses extends Responses {
                 latency: 0,
                 baseURL: this.baseURL ?? '',
                 params: body,
-                httpStatus: error?.status ? error.status : 500,
+                httpStatus,
                 usage: {
                   inputTokens: 0,
                   outputTokens: 0
@@ -839,6 +911,7 @@ let WrappedResponses$1 = class WrappedResponses extends Responses {
         }
         return result;
       }, async error => {
+        const httpStatus = error && typeof error === 'object' && 'status' in error ? error.status ?? 500 : 500;
         await sendEventToPosthog({
           client: this.phClient,
           distinctId: posthogDistinctId,
@@ -851,7 +924,7 @@ let WrappedResponses$1 = class WrappedResponses extends Responses {
           latency: 0,
           baseURL: this.baseURL ?? '',
           params: body,
-          httpStatus: error?.status ? error.status : 500,
+          httpStatus,
           usage: {
             inputTokens: 0,
             outputTokens: 0
@@ -910,6 +983,7 @@ let WrappedResponses$1 = class WrappedResponses extends Responses {
         });
         return result;
       }, async error => {
+        const httpStatus = error && typeof error === 'object' && 'status' in error ? error.status ?? 500 : 500;
         await sendEventToPosthog({
           client: this.phClient,
           distinctId: posthogDistinctId,
@@ -922,7 +996,7 @@ let WrappedResponses$1 = class WrappedResponses extends Responses {
           latency: 0,
           baseURL: this.baseURL ?? '',
           params: body,
-          httpStatus: error?.status ? error.status : 500,
+          httpStatus,
           usage: {
             inputTokens: 0,
             outputTokens: 0
@@ -984,14 +1058,52 @@ class WrappedCompletions extends AzureOpenAI.Chat.Completions {
           const [stream1, stream2] = value.tee();
           (async () => {
             try {
+              const contentBlocks = [];
               let accumulatedContent = '';
               let usage = {
                 inputTokens: 0,
                 outputTokens: 0
               };
+              // Map to track in-progress tool calls
+              const toolCallsInProgress = new Map();
               for await (const chunk of stream1) {
-                const delta = chunk?.choices?.[0]?.delta?.content ?? '';
-                accumulatedContent += delta;
+                const choice = chunk?.choices?.[0];
+                // Handle text content
+                const deltaContent = choice?.delta?.content;
+                if (deltaContent) {
+                  accumulatedContent += deltaContent;
+                }
+                // Handle tool calls
+                const deltaToolCalls = choice?.delta?.tool_calls;
+                if (deltaToolCalls && Array.isArray(deltaToolCalls)) {
+                  for (const toolCall of deltaToolCalls) {
+                    const index = toolCall.index;
+                    if (index !== undefined) {
+                      if (!toolCallsInProgress.has(index)) {
+                        // New tool call
+                        toolCallsInProgress.set(index, {
+                          id: toolCall.id || '',
+                          name: toolCall.function?.name || '',
+                          arguments: ''
+                        });
+                      }
+                      const inProgressCall = toolCallsInProgress.get(index);
+                      if (inProgressCall) {
+                        // Update tool call data
+                        if (toolCall.id) {
+                          inProgressCall.id = toolCall.id;
+                        }
+                        if (toolCall.function?.name) {
+                          inProgressCall.name = toolCall.function.name;
+                        }
+                        if (toolCall.function?.arguments) {
+                          inProgressCall.arguments += toolCall.function.arguments;
+                        }
+                      }
+                    }
+                  }
+                }
+                // Handle usage information
                 if (chunk.usage) {
                   usage = {
                     inputTokens: chunk.usage.prompt_tokens ?? 0,
@@ -1001,6 +1113,37 @@ class WrappedCompletions extends AzureOpenAI.Chat.Completions {
                   };
                 }
               }
+              // Build final content blocks
+              if (accumulatedContent) {
+                contentBlocks.push({
+                  type: 'text',
+                  text: accumulatedContent
+                });
+              }
+              // Add completed tool calls to content blocks
+              for (const toolCall of toolCallsInProgress.values()) {
+                if (toolCall.name) {
+                  contentBlocks.push({
+                    type: 'function',
+                    id: toolCall.id,
+                    function: {
+                      name: toolCall.name,
+                      arguments: toolCall.arguments
+                    }
+                  });
+                }
+              }
+              // Format output to match non-streaming version
+              const formattedOutput = contentBlocks.length > 0 ? [{
+                role: 'assistant',
+                content: contentBlocks
+              }] : [{
+                role: 'assistant',
+                content: [{
+                  type: 'text',
+                  text: ''
+                }]
+              }];
               const latency = (Date.now() - startTime) / 1000;
               await sendEventToPosthog({
                 client: this.phClient,
@@ -1009,10 +1152,7 @@ class WrappedCompletions extends AzureOpenAI.Chat.Completions {
                 model: openAIParams.model,
                 provider: 'azure',
                 input: openAIParams.messages,
-                output: [{
-                  content: accumulatedContent,
-                  role: 'assistant'
-                }],
+                output: formattedOutput,
                 latency,
                 baseURL: this.baseURL ?? '',
                 params: body,
@@ -1021,6 +1161,7 @@ class WrappedCompletions extends AzureOpenAI.Chat.Completions {
                 captureImmediate: posthogCaptureImmediate
               });
             } catch (error) {
+              const httpStatus = error && typeof error === 'object' && 'status' in error ? error.status ?? 500 : 500;
               await sendEventToPosthog({
                 client: this.phClient,
                 distinctId: posthogDistinctId,
@@ -1032,7 +1173,7 @@ class WrappedCompletions extends AzureOpenAI.Chat.Completions {
                 latency: 0,
                 baseURL: this.baseURL ?? '',
                 params: body,
-                httpStatus: error?.status ? error.status : 500,
+                httpStatus,
                 usage: {
                   inputTokens: 0,
                   outputTokens: 0
@@ -1075,6 +1216,7 @@ class WrappedCompletions extends AzureOpenAI.Chat.Completions {
         }
         return result;
       }, async error => {
+        const httpStatus = error && typeof error === 'object' && 'status' in error ? error.status ?? 500 : 500;
         await sendEventToPosthog({
           client: this.phClient,
           distinctId: posthogDistinctId,
@@ -1086,7 +1228,7 @@ class WrappedCompletions extends AzureOpenAI.Chat.Completions {
           latency: 0,
           baseURL: this.baseURL ?? '',
           params: body,
-          httpStatus: error?.status ? error.status : 500,
+          httpStatus,
           usage: {
             inputTokens: 0,
             outputTokens: 0
@@ -1163,6 +1305,7 @@ class WrappedResponses extends AzureOpenAI.Responses {
                 captureImmediate: posthogCaptureImmediate
               });
             } catch (error) {
+              const httpStatus = error && typeof error === 'object' && 'status' in error ? error.status ?? 500 : 500;
               await sendEventToPosthog({
                 client: this.phClient,
                 distinctId: posthogDistinctId,
@@ -1175,7 +1318,7 @@ class WrappedResponses extends AzureOpenAI.Responses {
                 latency: 0,
                 baseURL: this.baseURL ?? '',
                 params: body,
-                httpStatus: error?.status ? error.status : 500,
+                httpStatus,
                 usage: {
                   inputTokens: 0,
                   outputTokens: 0
@@ -1218,6 +1361,7 @@ class WrappedResponses extends AzureOpenAI.Responses {
         }
         return result;
       }, async error => {
+        const httpStatus = error && typeof error === 'object' && 'status' in error ? error.status ?? 500 : 500;
         await sendEventToPosthog({
           client: this.phClient,
           distinctId: posthogDistinctId,
@@ -1230,7 +1374,7 @@ class WrappedResponses extends AzureOpenAI.Responses {
           latency: 0,
           baseURL: this.baseURL ?? '',
           params: body,
-          httpStatus: error?.status ? error.status : 500,
+          httpStatus,
           usage: {
             inputTokens: 0,
             outputTokens: 0
@@ -1597,6 +1741,8 @@ const createInstrumentationMiddleware = (phClient, model, options) => {
       const provider = options.posthogProviderOverride ?? extractProvider(model);
       const availableTools = extractAvailableToolCalls('vercel', params);
       const baseURL = ''; // cannot currently get baseURL from vercel
+      // Map to track in-progress tool calls
+      const toolCallsInProgress = new Map();
       try {
         const {
           stream,
@@ -1611,6 +1757,34 @@ const createInstrumentationMiddleware = (phClient, model, options) => {
             if (chunk.type === 'reasoning-delta') {
               reasoningText += chunk.delta; // New in v5
             }
+            // Handle tool call chunks
+            if (chunk.type === 'tool-input-start') {
+              // Initialize a new tool call
+              toolCallsInProgress.set(chunk.id, {
+                toolCallId: chunk.id,
+                toolName: chunk.toolName,
+                input: ''
+              });
+            }
+            if (chunk.type === 'tool-input-delta') {
+              // Accumulate tool call arguments
+              const toolCall = toolCallsInProgress.get(chunk.id);
+              if (toolCall) {
+                toolCall.input += chunk.delta;
+              }
+            }
+            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)
+              toolCallsInProgress.set(chunk.toolCallId, {
+                toolCallId: chunk.toolCallId,
+                toolName: chunk.toolName,
+                input: chunk.input
+              });
+            }
             if (chunk.type === 'finish') {
               const providerMetadata = chunk.providerMetadata;
               const additionalTokenValues = {
@@ -1644,6 +1818,19 @@ const createInstrumentationMiddleware = (phClient, model, options) => {
                 text: truncate(generatedText)
               });
             }
+            // Add completed tool calls to content
+            for (const toolCall of toolCallsInProgress.values()) {
+              if (toolCall.toolName) {
+                content.push({
+                  type: 'tool-call',
+                  id: toolCall.toolCallId,
+                  function: {
+                    name: toolCall.toolName,
+                    arguments: toolCall.input
+                  }
+                });
+              }
+            }
             // Structure output like mapVercelOutput does
             const output = content.length > 0 ? [{
               role: 'assistant',
@@ -1746,6 +1933,9 @@ class WrappedMessages extends AnthropicOriginal.Messages {
     if (anthropicParams.stream) {
       return parentPromise.then(value => {
         let accumulatedContent = '';
+        const contentBlocks = [];
+        const toolsInProgress = new Map();
+        let currentTextBlock = null;
         const usage = {
           inputTokens: 0,
           outputTokens: 0,
@@ -1757,10 +1947,70 @@ class WrappedMessages extends AnthropicOriginal.Messages {
           (async () => {
             try {
               for await (const chunk of stream1) {
+                // Handle content block start events
+                if (chunk.type === 'content_block_start') {
+                  if (chunk.content_block?.type === 'text') {
+                    currentTextBlock = {
+                      type: 'text',
+                      text: ''
+                    };
+                    contentBlocks.push(currentTextBlock);
+                  } else if (chunk.content_block?.type === 'tool_use') {
+                    const toolBlock = {
+                      type: 'function',
+                      id: chunk.content_block.id,
+                      function: {
+                        name: chunk.content_block.name,
+                        arguments: {}
+                      }
+                    };
+                    contentBlocks.push(toolBlock);
+                    toolsInProgress.set(chunk.content_block.id, {
+                      block: toolBlock,
+                      inputString: ''
+                    });
+                    currentTextBlock = null;
+                  }
+                }
+                // Handle text delta events
                 if ('delta' in chunk) {
                   if ('text' in chunk.delta) {
                     const delta = chunk?.delta?.text ?? '';
                     accumulatedContent += delta;
+                    if (currentTextBlock) {
+                      currentTextBlock.text += delta;
+                    }
+                  }
+                }
+                // Handle tool input delta events
+                if (chunk.type === 'content_block_delta' && chunk.delta?.type === 'input_json_delta') {
+                  const block = chunk.index !== undefined ? contentBlocks[chunk.index] : undefined;
+                  const toolId = block?.type === 'function' ? block.id : undefined;
+                  if (toolId && toolsInProgress.has(toolId)) {
+                    const tool = toolsInProgress.get(toolId);
+                    if (tool) {
+                      tool.inputString += chunk.delta.partial_json || '';
+                    }
+                  }
+                }
+                // Handle content block stop events
+                if (chunk.type === 'content_block_stop') {
+                  currentTextBlock = null;
+                  // Parse accumulated tool input
+                  if (chunk.index !== undefined) {
+                    const block = contentBlocks[chunk.index];
+                    if (block?.type === 'function' && block.id && toolsInProgress.has(block.id)) {
+                      const tool = toolsInProgress.get(block.id);
+                      if (tool) {
+                        try {
+                          block.function.arguments = JSON.parse(tool.inputString);
+                        } catch (e) {
+                          // Keep empty object if parsing fails
+                          console.error('Error parsing tool input:', e);
+                        }
+                      }
+                      toolsInProgress.delete(block.id);
+                    }
                   }
                 }
                 if (chunk.type == 'message_start') {
@@ -1774,6 +2024,17 @@ class WrappedMessages extends AnthropicOriginal.Messages {
               }
               const latency = (Date.now() - startTime) / 1000;
               const availableTools = extractAvailableToolCalls('anthropic', anthropicParams);
+              // Format output to match non-streaming version
+              const formattedOutput = contentBlocks.length > 0 ? [{
+                role: 'assistant',
+                content: contentBlocks
+              }] : [{
+                role: 'assistant',
+                content: [{
+                  type: 'text',
+                  text: accumulatedContent
+                }]
+              }];
               await sendEventToPosthog({
                 client: this.phClient,
                 distinctId: posthogDistinctId,
@@ -1781,10 +2042,7 @@ class WrappedMessages extends AnthropicOriginal.Messages {
                 model: anthropicParams.model,
                 provider: 'anthropic',
                 input: sanitizeAnthropic(mergeSystemPrompt(anthropicParams, 'anthropic')),
-                output: [{
-                  content: accumulatedContent,
-                  role: 'assistant'
-                }],
+                output: formattedOutput,
                 latency,
                 baseURL: this.baseURL ?? '',
                 params: body,
@@ -1909,6 +2167,7 @@ class WrappedModels {
       const response = await this.client.models.generateContent(geminiParams);
       const latency = (Date.now() - startTime) / 1000;
       const availableTools = extractAvailableToolCalls('gemini', geminiParams);
+      const metadata = response.usageMetadata;
       await sendEventToPosthog({
         client: this.phClient,
         distinctId: posthogDistinctId,
@@ -1922,10 +2181,10 @@ class WrappedModels {
         params: params,
         httpStatus: 200,
         usage: {
-          inputTokens: response.usageMetadata?.promptTokenCount ?? 0,
-          outputTokens: response.usageMetadata?.candidatesTokenCount ?? 0,
-          reasoningTokens: response.usageMetadata?.thoughtsTokenCount ?? 0,
-          cacheReadInputTokens: response.usageMetadata?.cachedContentTokenCount ?? 0
+          inputTokens: metadata?.promptTokenCount ?? 0,
+          outputTokens: metadata?.candidatesTokenCount ?? 0,
+          reasoningTokens: metadata?.thoughtsTokenCount ?? 0,
+          cacheReadInputTokens: metadata?.cachedContentTokenCount ?? 0
         },
         tools: availableTools,
         captureImmediate: posthogCaptureImmediate
@@ -1967,7 +2226,7 @@ class WrappedModels {
     } = params;
     const traceId = posthogTraceId ?? v4();
     const startTime = Date.now();
-    let accumulatedContent = '';
+    const accumulatedContent = [];
     let usage = {
       inputTokens: 0,
       outputTokens: 0
@@ -1975,21 +2234,66 @@ class WrappedModels {
     try {
       const stream = await this.client.models.generateContentStream(geminiParams);
       for await (const chunk of stream) {
+        // Handle text content
         if (chunk.text) {
-          accumulatedContent += chunk.text;
+          // Find if we already have a text item to append to
+          let lastTextItem;
+          for (let i = accumulatedContent.length - 1; i >= 0; i--) {
+            if (accumulatedContent[i].type === 'text') {
+              lastTextItem = accumulatedContent[i];
+              break;
+            }
+          }
+          if (lastTextItem && lastTextItem.type === 'text') {
+            lastTextItem.text += chunk.text;
+          } else {
+            accumulatedContent.push({
+              type: 'text',
+              text: chunk.text
+            });
+          }
         }
+        // Handle function calls from candidates
+        if (chunk.candidates && Array.isArray(chunk.candidates)) {
+          for (const candidate of chunk.candidates) {
+            if (candidate.content && candidate.content.parts) {
+              for (const part of candidate.content.parts) {
+                // Type-safe check for functionCall
+                if ('functionCall' in part) {
+                  const funcCall = part.functionCall;
+                  if (funcCall?.name) {
+                    accumulatedContent.push({
+                      type: 'function',
+                      function: {
+                        name: funcCall.name,
+                        arguments: funcCall.args || {}
+                      }
+                    });
+                  }
+                }
+              }
+            }
+          }
+        }
+        // Update usage metadata - handle both old and new field names
         if (chunk.usageMetadata) {
+          const metadata = chunk.usageMetadata;
           usage = {
-            inputTokens: chunk.usageMetadata.promptTokenCount ?? 0,
-            outputTokens: chunk.usageMetadata.candidatesTokenCount ?? 0,
-            reasoningTokens: chunk.usageMetadata.thoughtsTokenCount ?? 0,
-            cacheReadInputTokens: chunk.usageMetadata.cachedContentTokenCount ?? 0
+            inputTokens: metadata.promptTokenCount ?? 0,
+            outputTokens: metadata.candidatesTokenCount ?? 0,
+            reasoningTokens: metadata.thoughtsTokenCount ?? 0,
+            cacheReadInputTokens: metadata.cachedContentTokenCount ?? 0
           };
         }
         yield chunk;
       }
       const latency = (Date.now() - startTime) / 1000;
       const availableTools = extractAvailableToolCalls('gemini', geminiParams);
+      // Format output similar to formatResponseGemini
+      const output = accumulatedContent.length > 0 ? [{
+        role: 'assistant',
+        content: accumulatedContent
+      }] : [];
       await sendEventToPosthog({
         client: this.phClient,
         distinctId: posthogDistinctId,
@@ -1997,10 +2301,7 @@ class WrappedModels {
         model: geminiParams.model,
         provider: 'gemini',
         input: this.formatInputForPostHog(geminiParams.contents),
-        output: [{
-          content: accumulatedContent,
-          role: 'assistant'
-        }],
+        output,
         latency,
         baseURL: 'https://generativelanguage.googleapis.com',
         params: params,
@@ -2050,22 +2351,28 @@ class WrappedModels {
           };
         }
         if (item && typeof item === 'object') {
-          if (item.text) {
+          const obj = item;
+          if ('text' in obj && obj.text) {
             return {
-              role: item.role || 'user',
-              content: item.text
+              role: obj.role || 'user',
+              content: obj.text
             };
           }
-          if (item.content) {
+          if ('content' in obj && obj.content) {
             return {
-              role: item.role || 'user',
-              content: item.content
+              role: obj.role || 'user',
+              content: obj.content
             };
           }
-          if (item.parts) {
+          if ('parts' in obj && Array.isArray(obj.parts)) {
             return {
-              role: item.role || 'user',
-              content: item.parts.map(part => part.text ? part.text : part)
+              role: obj.role || 'user',
+              content: obj.parts.map(part => {
+                if (part && typeof part === 'object' && 'text' in part) {
+                  return part.text;
+                }
+                return part;
+              })
             };
           }
         }
@@ -2076,16 +2383,17 @@ class WrappedModels {
       });
     }
     if (contents && typeof contents === 'object') {
-      if (contents.text) {
+      const obj = contents;
+      if ('text' in obj && obj.text) {
         return [{
           role: 'user',
-          content: contents.text
+          content: obj.text
         }];
       }
-      if (contents.content) {
+      if ('content' in obj && obj.content) {
         return [{
           role: 'user',
-          content: contents.content
+          content: obj.content
         }];
       }
     }
@@ -2850,6 +3158,8 @@ class LangChainCallbackHandler extends BaseCallbackHandler {
     const eventName = parentRunId ? '$ai_span' : '$ai_trace';
     const latency = run.endTime ? (run.endTime - run.startTime) / 1000 : 0;
     const eventProperties = {
+      $ai_lib: 'posthog-ai',
+      $ai_lib_version: version,
       $ai_trace_id: traceId,
       $ai_input_state: withPrivacyMode(this.client, this.privacyMode, run.input),
       $ai_latency: latency,
@@ -2890,6 +3200,8 @@ class LangChainCallbackHandler extends BaseCallbackHandler {
   _captureGeneration(traceId, runId, run, output, parentRunId) {
     const latency = run.endTime ? (run.endTime - run.startTime) / 1000 : 0;
     const eventProperties = {
+      $ai_lib: 'posthog-ai',
+      $ai_lib_version: version,
       $ai_trace_id: traceId,
       $ai_span_id: runId,
       $ai_span_name: run.name,