@posthog/ai
Version:
PostHog Node.js AI integrations
613 lines (596 loc) • 19.2 kB
JavaScript
import { wrapLanguageModel } from 'ai';
import { v4 } from 'uuid';
import { Buffer } from 'buffer';
// limit large outputs by truncating to 200kb (approx 200k bytes)
const MAX_OUTPUT_SIZE = 200000;
const STRING_FORMAT = 'utf8';
const getModelParams = params => {
if (!params) {
return {};
}
const modelParams = {};
const paramKeys = ['temperature', 'max_tokens', 'max_completion_tokens', 'top_p', 'frequency_penalty', 'presence_penalty', 'n', 'stop', 'stream', 'streaming'];
for (const key of paramKeys) {
if (key in params && params[key] !== undefined) {
modelParams[key] = params[key];
}
}
return modelParams;
};
const withPrivacyMode = (client, privacyMode, input) => {
return client.privacy_mode || privacyMode ? null : input;
};
const truncate = str => {
try {
const buffer = Buffer.from(str, STRING_FORMAT);
if (buffer.length <= MAX_OUTPUT_SIZE) {
return str;
}
const truncatedBuffer = buffer.slice(0, MAX_OUTPUT_SIZE);
return `${truncatedBuffer.toString(STRING_FORMAT)}... [truncated]`;
} catch (error) {
console.error('Error truncating, likely not a string');
return str;
}
};
/**
* Extract available tool calls from the request parameters.
* These are the tools provided to the LLM, not the tool calls in the response.
*/
const extractAvailableToolCalls = (provider, params) => {
{
// 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;
}
return null;
}
};
function sanitizeValues(obj) {
if (obj === undefined || obj === null) {
return obj;
}
const jsonSafe = JSON.parse(JSON.stringify(obj));
if (typeof jsonSafe === 'string') {
return Buffer.from(jsonSafe, STRING_FORMAT).toString(STRING_FORMAT);
} else if (Array.isArray(jsonSafe)) {
return jsonSafe.map(sanitizeValues);
} else if (jsonSafe && typeof jsonSafe === 'object') {
return Object.fromEntries(Object.entries(jsonSafe).map(([k, v]) => [k, sanitizeValues(v)]));
}
return jsonSafe;
}
const sendEventToPosthog = async ({
client,
distinctId,
traceId,
model,
provider,
input,
output,
latency,
baseURL,
params,
httpStatus = 200,
usage = {},
isError = false,
error,
tools,
captureImmediate = false
}) => {
if (!client.capture) {
return Promise.resolve();
}
// sanitize input and output for UTF-8 validity
const safeInput = sanitizeValues(input);
const safeOutput = sanitizeValues(output);
const safeError = sanitizeValues(error);
let errorData = {};
if (isError) {
errorData = {
$ai_is_error: true,
$ai_error: safeError
};
}
let costOverrideData = {};
if (params.posthogCostOverride) {
const inputCostUSD = (params.posthogCostOverride.inputCost ?? 0) * (usage.inputTokens ?? 0);
const outputCostUSD = (params.posthogCostOverride.outputCost ?? 0) * (usage.outputTokens ?? 0);
costOverrideData = {
$ai_input_cost_usd: inputCostUSD,
$ai_output_cost_usd: outputCostUSD,
$ai_total_cost_usd: inputCostUSD + outputCostUSD
};
}
const additionalTokenValues = {
...(usage.reasoningTokens ? {
$ai_reasoning_tokens: usage.reasoningTokens
} : {}),
...(usage.cacheReadInputTokens ? {
$ai_cache_read_input_tokens: usage.cacheReadInputTokens
} : {}),
...(usage.cacheCreationInputTokens ? {
$ai_cache_creation_input_tokens: usage.cacheCreationInputTokens
} : {})
};
const properties = {
$ai_provider: params.posthogProviderOverride ?? provider,
$ai_model: params.posthogModelOverride ?? model,
$ai_model_parameters: getModelParams(params),
$ai_input: withPrivacyMode(client, params.posthogPrivacyMode ?? false, safeInput),
$ai_output_choices: withPrivacyMode(client, params.posthogPrivacyMode ?? false, safeOutput),
$ai_http_status: httpStatus,
$ai_input_tokens: usage.inputTokens ?? 0,
$ai_output_tokens: usage.outputTokens ?? 0,
...additionalTokenValues,
$ai_latency: latency,
$ai_trace_id: traceId,
$ai_base_url: baseURL,
...params.posthogProperties,
...(distinctId ? {} : {
$process_person_profile: false
}),
...(tools ? {
$ai_tools: tools
} : {}),
...errorData,
...costOverrideData
};
const event = {
distinctId: distinctId ?? traceId,
event: '$ai_generation',
properties,
groups: params.posthogGroups
};
if (captureImmediate) {
// await capture promise to send single event in serverless environments
await client.captureImmediate(event);
} else {
client.capture(event);
}
};
// Type guards for safer type checking
const isString = value => {
return typeof value === 'string';
};
const REDACTED_IMAGE_PLACEHOLDER = '[base64 image redacted]';
// ============================================
// Base64 Detection Helpers
// ============================================
const isBase64DataUrl = str => {
return /^data:([^;]+);base64,/.test(str);
};
const isValidUrl = str => {
try {
new URL(str);
return true;
} catch {
// Not an absolute URL, check if it's a relative URL or path
return str.startsWith('/') || str.startsWith('./') || str.startsWith('../');
}
};
const isRawBase64 = str => {
// Skip if it's a valid URL or path
if (isValidUrl(str)) {
return false;
}
// Check if it's a valid base64 string
// Base64 images are typically at least a few hundred chars, but we'll be conservative
return str.length > 20 && /^[A-Za-z0-9+/]+=*$/.test(str);
};
function redactBase64DataUrl(str) {
if (!isString(str)) return str;
// Check for data URL format
if (isBase64DataUrl(str)) {
return REDACTED_IMAGE_PLACEHOLDER;
}
// Check for raw base64 (Vercel sends raw base64 for inline images)
if (isRawBase64(str)) {
return REDACTED_IMAGE_PLACEHOLDER;
}
return str;
}
const mapVercelParams = params => {
return {
temperature: params.temperature,
max_output_tokens: params.maxOutputTokens,
top_p: params.topP,
frequency_penalty: params.frequencyPenalty,
presence_penalty: params.presencePenalty,
stop: params.stopSequences,
stream: params.stream
};
};
const mapVercelPrompt = messages => {
// Map and truncate individual content
const inputs = messages.map(message => {
let content;
// Handle system role which has string content
if (message.role === 'system') {
content = [{
type: 'text',
text: truncate(String(message.content))
}];
} else {
// Handle other roles which have array content
if (Array.isArray(message.content)) {
content = message.content.map(c => {
if (c.type === 'text') {
return {
type: 'text',
text: truncate(c.text)
};
} else if (c.type === 'file') {
// For file type, check if it's a data URL and redact if needed
let fileData;
const contentData = c.data;
if (contentData instanceof URL) {
fileData = contentData.toString();
} else if (isString(contentData)) {
// Redact base64 data URLs and raw base64 to prevent oversized events
fileData = redactBase64DataUrl(contentData);
} else {
fileData = 'raw files not supported';
}
return {
type: 'file',
file: fileData,
mediaType: c.mediaType
};
} else if (c.type === 'reasoning') {
return {
type: 'reasoning',
text: truncate(c.reasoning)
};
} else if (c.type === 'tool-call') {
return {
type: 'tool-call',
toolCallId: c.toolCallId,
toolName: c.toolName,
input: c.input
};
} else if (c.type === 'tool-result') {
return {
type: 'tool-result',
toolCallId: c.toolCallId,
toolName: c.toolName,
output: c.output,
isError: c.isError
};
}
return {
type: 'text',
text: ''
};
});
} else {
// Fallback for non-array content
content = [{
type: 'text',
text: truncate(String(message.content))
}];
}
}
return {
role: message.role,
content
};
});
try {
// Trim the inputs array until its JSON size fits within MAX_OUTPUT_SIZE
let serialized = JSON.stringify(inputs);
let removedCount = 0;
// We need to keep track of the initial size of the inputs array because we're going to be mutating it
const initialSize = inputs.length;
for (let i = 0; i < initialSize && Buffer.byteLength(serialized, 'utf8') > MAX_OUTPUT_SIZE; i++) {
inputs.shift();
removedCount++;
serialized = JSON.stringify(inputs);
}
if (removedCount > 0) {
// Add one placeholder to indicate how many were removed
inputs.unshift({
role: 'posthog',
content: `[${removedCount} message${removedCount === 1 ? '' : 's'} removed due to size limit]`
});
}
} catch (error) {
console.error('Error stringifying inputs', error);
return [{
role: 'posthog',
content: 'An error occurred while processing your request. Please try again.'
}];
}
return inputs;
};
const mapVercelOutput = result => {
const content = result.map(item => {
if (item.type === 'text') {
return {
type: 'text',
text: truncate(item.text)
};
}
if (item.type === 'tool-call') {
return {
type: 'tool-call',
id: item.toolCallId,
function: {
name: item.toolName,
arguments: item.args || JSON.stringify(item.arguments || {})
}
};
}
if (item.type === 'reasoning') {
return {
type: 'reasoning',
text: truncate(item.text)
};
}
if (item.type === 'file') {
// Handle files similar to input mapping - avoid large base64 data
let fileData;
if (item.data instanceof URL) {
fileData = item.data.toString();
} else if (typeof item.data === 'string') {
fileData = redactBase64DataUrl(item.data);
// If not redacted and still large, replace with size indicator
if (fileData === item.data && item.data.length > 1000) {
fileData = `[${item.mediaType} file - ${item.data.length} bytes]`;
}
} else {
fileData = `[binary ${item.mediaType} file]`;
}
return {
type: 'file',
name: 'generated_file',
mediaType: item.mediaType,
data: fileData
};
}
if (item.type === 'source') {
return {
type: 'source',
sourceType: item.sourceType,
id: item.id,
url: item.url || '',
title: item.title || ''
};
}
// Fallback for unknown types - try to extract text if possible
return {
type: 'text',
text: truncate(JSON.stringify(item))
};
});
if (content.length > 0) {
return [{
role: 'assistant',
content: content.length === 1 && content[0].type === 'text' ? content[0].text : content
}];
}
// otherwise stringify and truncate
try {
const jsonOutput = JSON.stringify(result);
return [{
content: truncate(jsonOutput),
role: 'assistant'
}];
} catch (error) {
console.error('Error stringifying output');
return [];
}
};
const extractProvider = model => {
const provider = model.provider.toLowerCase();
const providerName = provider.split('.')[0];
return providerName;
};
const createInstrumentationMiddleware = (phClient, model, options) => {
const middleware = {
wrapGenerate: async ({
doGenerate,
params
}) => {
const startTime = Date.now();
const mergedParams = {
...options,
...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 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 usage = {
inputTokens: result.usage.inputTokens,
outputTokens: result.usage.outputTokens,
reasoningTokens: result.usage.reasoningTokens,
cacheReadInputTokens: result.usage.cachedInputTokens,
...additionalTokenValues
};
await sendEventToPosthog({
client: phClient,
distinctId: options.posthogDistinctId,
traceId: options.posthogTraceId ?? v4(),
model: modelId,
provider: provider,
input: options.posthogPrivacyMode ? '' : mapVercelPrompt(params.prompt),
output: content,
latency,
baseURL,
params: mergedParams,
httpStatus: 200,
usage,
tools: availableTools,
captureImmediate: options.posthogCaptureImmediate
});
return result;
} catch (error) {
const modelId = model.modelId;
await sendEventToPosthog({
client: phClient,
distinctId: options.posthogDistinctId,
traceId: options.posthogTraceId ?? v4(),
model: modelId,
provider: model.provider,
input: options.posthogPrivacyMode ? '' : mapVercelPrompt(params.prompt),
output: [],
latency: 0,
baseURL: '',
params: mergedParams,
httpStatus: error?.status ? error.status : 500,
usage: {
inputTokens: 0,
outputTokens: 0
},
isError: true,
error: truncate(JSON.stringify(error)),
tools: availableTools,
captureImmediate: options.posthogCaptureImmediate
});
throw error;
}
},
wrapStream: async ({
doStream,
params
}) => {
const startTime = Date.now();
let generatedText = '';
let reasoningText = '';
let usage = {};
const mergedParams = {
...options,
...mapVercelParams(params)
};
const modelId = options.posthogModelOverride ?? model.modelId;
const provider = options.posthogProviderOverride ?? extractProvider(model);
const availableTools = extractAvailableToolCalls('vercel', params);
const baseURL = ''; // cannot currently get baseURL from vercel
try {
const {
stream,
...rest
} = await doStream();
const transformStream = new TransformStream({
transform(chunk, controller) {
// Handle new v5 streaming patterns
if (chunk.type === 'text-delta') {
generatedText += chunk.delta;
}
if (chunk.type === 'reasoning-delta') {
reasoningText += chunk.delta; // New in v5
}
if (chunk.type === 'finish') {
const providerMetadata = chunk.providerMetadata;
const additionalTokenValues = {
...(providerMetadata?.anthropic ? {
cacheCreationInputTokens: providerMetadata.anthropic.cacheCreationInputTokens
} : {})
};
usage = {
inputTokens: chunk.usage?.inputTokens,
outputTokens: chunk.usage?.outputTokens,
reasoningTokens: chunk.usage?.reasoningTokens,
cacheReadInputTokens: chunk.usage?.cachedInputTokens,
...additionalTokenValues
};
}
controller.enqueue(chunk);
},
flush: async () => {
const latency = (Date.now() - startTime) / 1000;
// Build content array similar to mapVercelOutput structure
const content = [];
if (reasoningText) {
content.push({
type: 'reasoning',
text: truncate(reasoningText)
});
}
if (generatedText) {
content.push({
type: 'text',
text: truncate(generatedText)
});
}
// Structure output like mapVercelOutput does
const output = content.length > 0 ? [{
role: 'assistant',
content: content.length === 1 && content[0].type === 'text' ? content[0].text : content
}] : [];
await sendEventToPosthog({
client: phClient,
distinctId: options.posthogDistinctId,
traceId: options.posthogTraceId ?? v4(),
model: modelId,
provider: provider,
input: options.posthogPrivacyMode ? '' : mapVercelPrompt(params.prompt),
output: output,
latency,
baseURL,
params: mergedParams,
httpStatus: 200,
usage,
tools: availableTools,
captureImmediate: options.posthogCaptureImmediate
});
}
});
return {
stream: stream.pipeThrough(transformStream),
...rest
};
} catch (error) {
await sendEventToPosthog({
client: phClient,
distinctId: options.posthogDistinctId,
traceId: options.posthogTraceId ?? v4(),
model: modelId,
provider: provider,
input: options.posthogPrivacyMode ? '' : mapVercelPrompt(params.prompt),
output: [],
latency: 0,
baseURL: '',
params: mergedParams,
httpStatus: error?.status ? error.status : 500,
usage: {
inputTokens: 0,
outputTokens: 0
},
isError: true,
error: truncate(JSON.stringify(error)),
tools: availableTools,
captureImmediate: options.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;
};
export { wrapVercelLanguageModel as withTracing };
//# sourceMappingURL=index.mjs.map