UNPKG

@posthog/ai

Version:
613 lines (596 loc) 19.2 kB
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