diff --git a/packages/core/src/shared-exports.ts b/packages/core/src/shared-exports.ts index ffd08180b202..d9891976b597 100644 --- a/packages/core/src/shared-exports.ts +++ b/packages/core/src/shared-exports.ts @@ -221,6 +221,8 @@ export type { LangChainOptions, LangChainIntegration } from './tracing/langchain export { instrumentStateGraphCompile, instrumentCreateReactAgent, instrumentLangGraph } from './tracing/langgraph'; export { LANGGRAPH_INTEGRATION_NAME } from './tracing/langgraph/constants'; export type { LangGraphOptions, LangGraphIntegration, CompiledGraph } from './tracing/langgraph/types'; +export { instrumentWorkersAiClient } from './tracing/workers-ai'; +export type { WorkersAiClient, WorkersAiOptions } from './tracing/workers-ai/types'; // eslint-disable-next-line typescript/no-deprecated export type { OpenAiClient, OpenAiOptions, InstrumentedMethod } from './tracing/openai/types'; export type { diff --git a/packages/core/src/tracing/workers-ai/constants.ts b/packages/core/src/tracing/workers-ai/constants.ts new file mode 100644 index 000000000000..89847abbae1c --- /dev/null +++ b/packages/core/src/tracing/workers-ai/constants.ts @@ -0,0 +1,10 @@ +/** + * The provider value for the `gen_ai.system` attribute. + * @see https://developers.cloudflare.com/workers-ai/ + */ +export const WORKERS_AI_SYSTEM_NAME = 'cloudflare.workers_ai'; + +/** + * The Sentry origin for spans created by the Workers AI instrumentation. + */ +export const WORKERS_AI_ORIGIN = 'auto.ai.cloudflare.workers_ai'; diff --git a/packages/core/src/tracing/workers-ai/index.ts b/packages/core/src/tracing/workers-ai/index.ts new file mode 100644 index 000000000000..e38d4bbd53ce --- /dev/null +++ b/packages/core/src/tracing/workers-ai/index.ts @@ -0,0 +1,140 @@ +import { captureException } from '../../exports'; +import { SPAN_STATUS_ERROR } from '../../tracing'; +import { startSpan, startSpanManual } from '../../tracing/trace'; +import type { Span } from '../../types/span'; +import { resolveAIRecordingOptions, shouldEnableTruncation } from '../ai/utils'; +import { WORKERS_AI_ORIGIN } from './constants'; +import { instrumentWorkersAiStream } from './streaming'; +import type { WorkersAiOptions } from './types'; +import { addRequestAttributes, addResponseAttributes, extractRequestAttributes, getOperationName } from './utils'; + +function isReadableStream(value: unknown): value is ReadableStream { + return typeof ReadableStream !== 'undefined' && value instanceof ReadableStream; +} + +/** + * Wrap the `run` method of the Workers AI binding with Sentry tracing. + */ +function instrumentRun( + originalRun: (...args: unknown[]) => Promise, + context: unknown, + options: WorkersAiOptions & Required>, +): (...args: unknown[]) => Promise { + return function instrumentedRun(...args: unknown[]): Promise { + const [model, inputs, runOptions] = args as [unknown, unknown, Record | undefined]; + + const operationName = getOperationName(inputs); + const requestAttributes = extractRequestAttributes(model, inputs, operationName); + const modelName = typeof model === 'string' ? model : 'unknown'; + + const isStreamRequested = + !!inputs && typeof inputs === 'object' && (inputs as { stream?: unknown }).stream === true; + const returnsRawResponse = + !!runOptions && + typeof runOptions === 'object' && + (runOptions.returnRawResponse === true || runOptions.websocket === true); + + const spanConfig = { + name: `${operationName} ${modelName}`, + op: `gen_ai.${operationName}`, + attributes: requestAttributes, + }; + + if (isStreamRequested && !returnsRawResponse) { + return startSpanManual(spanConfig, (span: Span) => { + const originalResult = originalRun.apply(context, args) as Promise; + + if (options.recordInputs) { + addRequestAttributes(span, inputs, operationName, shouldEnableTruncation(options.enableTruncation)); + } + + return originalResult.then( + result => { + if (isReadableStream(result)) { + return instrumentWorkersAiStream(result, span, options.recordOutputs); + } + + // The model did not actually return a stream — finalize the span eagerly. + addResponseAttributes(span, result, options.recordOutputs); + span.end(); + return result; + }, + error => { + span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + captureException(error, { + mechanism: { handled: false, type: `${WORKERS_AI_ORIGIN}.stream`, data: { function: 'run' } }, + }); + span.end(); + throw error; + }, + ); + }); + } + + return startSpan(spanConfig, (span: Span) => { + const originalResult = originalRun.apply(context, args) as Promise; + + if (options.recordInputs) { + addRequestAttributes(span, inputs, operationName, shouldEnableTruncation(options.enableTruncation)); + } + + return originalResult.then( + result => { + if (!returnsRawResponse) { + addResponseAttributes(span, result, options.recordOutputs); + } + return result; + }, + error => { + captureException(error, { + mechanism: { handled: false, type: WORKERS_AI_ORIGIN, data: { function: 'run' } }, + }); + throw error; + }, + ); + }); + }; +} + +// Guards against double-wrapping when a binding instrumented via the automatic +// env instrumentation in `@sentry/cloudflare` is additionally wrapped manually. +const instrumentedClients = new WeakSet(); + +/** + * Instrument a Cloudflare Workers AI binding (`env.AI`) with Sentry tracing. + * + * This wraps the binding's `run` method to create `gen_ai` spans following the + * Sentry AI Agents conventions. All other methods are passed through untouched. + * + * In `@sentry/cloudflare`, the `env.AI` binding is instrumented automatically — + * wrapping manually is only needed to pass custom options. + * + * @example + * ```javascript + * const ai = Sentry.instrumentWorkersAiClient(env.AI, { recordInputs: true, recordOutputs: true }); + * const result = await ai.run('@cf/meta/llama-3.1-8b-instruct', { prompt: 'Hello' }); + * ``` + */ +export function instrumentWorkersAiClient(client: T, options?: WorkersAiOptions): T { + if (instrumentedClients.has(client)) { + return client; + } + + const resolvedOptions = resolveAIRecordingOptions(options); + + const instrumented = new Proxy(client, { + get(target: object, prop: string | symbol, receiver: unknown): unknown { + const value = Reflect.get(target, prop, receiver); + + if (prop === 'run' && typeof value === 'function') { + return instrumentRun(value as (...args: unknown[]) => Promise, target, resolvedOptions); + } + + // Bind passed-through functions to the original target to preserve `this` (e.g. private fields). + return typeof value === 'function' ? (value as (...args: unknown[]) => unknown).bind(target) : value; + }, + }) as T; + + instrumentedClients.add(instrumented); + return instrumented; +} diff --git a/packages/core/src/tracing/workers-ai/streaming.ts b/packages/core/src/tracing/workers-ai/streaming.ts new file mode 100644 index 000000000000..c2a1055a8be0 --- /dev/null +++ b/packages/core/src/tracing/workers-ai/streaming.ts @@ -0,0 +1,131 @@ +import { captureException } from '../../exports'; +import { SPAN_STATUS_ERROR } from '../../tracing'; +import type { Span } from '../../types/span'; +import { endStreamSpan, type StreamResponseState } from '../ai/utils'; +import { WORKERS_AI_ORIGIN } from './constants'; +import type { WorkersAiUsage } from './types'; + +interface WorkersAiStreamChunk { + response?: unknown; + usage?: WorkersAiUsage; + tool_calls?: unknown[]; +} + +/** + * Parse a single SSE line (`data: {...}`) and accumulate its data into the streaming state. + */ +function processLine(line: string, state: StreamResponseState, recordOutputs: boolean): void { + const trimmed = line.trim(); + if (!trimmed.startsWith('data:')) { + return; + } + + const data = trimmed.slice('data:'.length).trim(); + if (!data || data === '[DONE]') { + return; + } + + let parsed: WorkersAiStreamChunk; + try { + parsed = JSON.parse(data) as WorkersAiStreamChunk; + } catch { + return; + } + + if (parsed.usage) { + if (typeof parsed.usage.prompt_tokens === 'number') { + state.promptTokens = parsed.usage.prompt_tokens; + } + if (typeof parsed.usage.completion_tokens === 'number') { + state.completionTokens = parsed.usage.completion_tokens; + } + if (typeof parsed.usage.total_tokens === 'number') { + state.totalTokens = parsed.usage.total_tokens; + } + } + + if (recordOutputs && typeof parsed.response === 'string') { + state.responseTexts.push(parsed.response); + } + + if (recordOutputs && Array.isArray(parsed.tool_calls) && parsed.tool_calls.length > 0) { + state.toolCalls.push(...parsed.tool_calls); + } +} + +/** + * Wrap a Workers AI streaming response (a server-sent-events `ReadableStream`) so we can + * accumulate the response text and token usage while passing the original bytes through untouched. + * + * The span is ended once the consumer finishes reading (or cancels) the stream. + */ +export function instrumentWorkersAiStream( + stream: ReadableStream, + span: Span, + recordOutputs: boolean, +): ReadableStream { + const reader = stream.getReader(); + const decoder = new TextDecoder(); + + const state: StreamResponseState = { + responseId: '', + responseModel: '', + finishReasons: [], + responseTexts: [], + toolCalls: [], + promptTokens: undefined, + completionTokens: undefined, + totalTokens: undefined, + }; + + let buffer = ''; + let spanEnded = false; + + const finish = (): void => { + if (spanEnded) { + return; + } + spanEnded = true; + endStreamSpan(span, state, recordOutputs); + }; + + const flushBuffer = (isDone: boolean): void => { + const lines = buffer.split('\n'); + // Keep the last (potentially incomplete) line in the buffer unless the stream is done. + buffer = isDone ? '' : (lines.pop() ?? ''); + for (const line of lines) { + processLine(line, state, recordOutputs); + } + }; + + return new ReadableStream({ + async pull(controller) { + try { + const { done, value } = await reader.read(); + + if (done) { + buffer += decoder.decode(); + flushBuffer(true); + finish(); + controller.close(); + return; + } + + buffer += decoder.decode(value, { stream: true }); + flushBuffer(false); + controller.enqueue(value); + } catch (error) { + span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + captureException(error, { + mechanism: { handled: false, type: `${WORKERS_AI_ORIGIN}.stream` }, + }); + finish(); + controller.error(error); + } + }, + async cancel(reason) { + finish(); + await reader.cancel(reason); + }, + }); +} diff --git a/packages/core/src/tracing/workers-ai/types.ts b/packages/core/src/tracing/workers-ai/types.ts new file mode 100644 index 000000000000..33334a6ed1bf --- /dev/null +++ b/packages/core/src/tracing/workers-ai/types.ts @@ -0,0 +1,58 @@ +import type { AIRecordingOptions } from '../ai/utils'; + +export interface WorkersAiOptions extends AIRecordingOptions { + /** + * Enable or disable truncation of recorded input messages. + * Defaults to `true`. + */ + enableTruncation?: boolean; +} + +/** + * Minimal shape of the Cloudflare Workers AI binding (`env.AI`). + * We only rely on the `run` method, everything else is passed through untouched. + * @see https://developers.cloudflare.com/workers-ai/configuration/bindings/ + */ +export interface WorkersAiClient { + run: (model: string, inputs: Record, options?: Record) => Promise; + [key: string]: unknown; +} + +/** + * The token usage reported by Workers AI text generation models. + */ +export interface WorkersAiUsage { + prompt_tokens?: number; + completion_tokens?: number; + total_tokens?: number; +} + +/** + * The (subset of) inputs accepted by Workers AI `run` calls that we read from. + * @see https://developers.cloudflare.com/workers-ai/models/ + */ +export interface WorkersAiInput { + // Text generation + prompt?: string; + messages?: Array<{ role?: string; content?: unknown }>; + stream?: boolean; + max_tokens?: number; + temperature?: number; + top_p?: number; + top_k?: number; + frequency_penalty?: number; + presence_penalty?: number; + tools?: unknown; + functions?: unknown; + // Text embeddings + text?: string | string[]; +} + +/** + * The (subset of) outputs returned by Workers AI text generation models that we read from. + */ +export interface WorkersAiOutput { + response?: unknown; + tool_calls?: unknown[]; + usage?: WorkersAiUsage; +} diff --git a/packages/core/src/tracing/workers-ai/utils.ts b/packages/core/src/tracing/workers-ai/utils.ts new file mode 100644 index 000000000000..94611c90ce19 --- /dev/null +++ b/packages/core/src/tracing/workers-ai/utils.ts @@ -0,0 +1,173 @@ +import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '../../semanticAttributes'; +import type { Span, SpanAttributeValue } from '../../types/span'; +import { + GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE, + GEN_AI_INPUT_MESSAGES_ATTRIBUTE, + GEN_AI_INPUT_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE, + GEN_AI_OPERATION_NAME_ATTRIBUTE, + GEN_AI_REQUEST_FREQUENCY_PENALTY_ATTRIBUTE, + GEN_AI_REQUEST_MAX_TOKENS_ATTRIBUTE, + GEN_AI_REQUEST_MODEL_ATTRIBUTE, + GEN_AI_REQUEST_PRESENCE_PENALTY_ATTRIBUTE, + GEN_AI_REQUEST_STREAM_ATTRIBUTE, + GEN_AI_REQUEST_TEMPERATURE_ATTRIBUTE, + GEN_AI_REQUEST_TOP_K_ATTRIBUTE, + GEN_AI_REQUEST_TOP_P_ATTRIBUTE, + GEN_AI_RESPONSE_TEXT_ATTRIBUTE, + GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE, + GEN_AI_SYSTEM_ATTRIBUTE, + GEN_AI_SYSTEM_INSTRUCTIONS_ATTRIBUTE, +} from '../ai/gen-ai-attributes'; +import { extractSystemInstructions, getJsonString, getTruncatedJsonString, setTokenUsageAttributes } from '../ai/utils'; +import { WORKERS_AI_ORIGIN, WORKERS_AI_SYSTEM_NAME } from './constants'; +import type { WorkersAiInput, WorkersAiOutput } from './types'; + +/** + * Determine the gen_ai operation name from the inputs passed to `AI.run`. + * Workers AI exposes a single `run` method, so we infer the operation from the input shape. + */ +export function getOperationName(inputs: unknown): string { + if (inputs && typeof inputs === 'object') { + if ('messages' in inputs || 'prompt' in inputs) { + return 'chat'; + } + if ('text' in inputs) { + return 'embeddings'; + } + } + return 'chat'; +} + +/** + * Extract the request attributes (model, request parameters, system, origin) from a `run` call. + */ +export function extractRequestAttributes( + model: unknown, + inputs: unknown, + operationName: string, +): Record { + const attributes: Record = { + [GEN_AI_SYSTEM_ATTRIBUTE]: WORKERS_AI_SYSTEM_NAME, + [GEN_AI_OPERATION_NAME_ATTRIBUTE]: operationName, + [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: WORKERS_AI_ORIGIN, + [GEN_AI_REQUEST_MODEL_ATTRIBUTE]: typeof model === 'string' ? model : 'unknown', + }; + + if (inputs && typeof inputs === 'object') { + const params = inputs as WorkersAiInput; + + if (typeof params.temperature === 'number') { + attributes[GEN_AI_REQUEST_TEMPERATURE_ATTRIBUTE] = params.temperature; + } + if (typeof params.max_tokens === 'number') { + attributes[GEN_AI_REQUEST_MAX_TOKENS_ATTRIBUTE] = params.max_tokens; + } + if (typeof params.top_p === 'number') { + attributes[GEN_AI_REQUEST_TOP_P_ATTRIBUTE] = params.top_p; + } + if (typeof params.top_k === 'number') { + attributes[GEN_AI_REQUEST_TOP_K_ATTRIBUTE] = params.top_k; + } + if (typeof params.frequency_penalty === 'number') { + attributes[GEN_AI_REQUEST_FREQUENCY_PENALTY_ATTRIBUTE] = params.frequency_penalty; + } + if (typeof params.presence_penalty === 'number') { + attributes[GEN_AI_REQUEST_PRESENCE_PENALTY_ATTRIBUTE] = params.presence_penalty; + } + if (params.stream === true) { + attributes[GEN_AI_REQUEST_STREAM_ATTRIBUTE] = true; + } + } + + return attributes; +} + +/** + * Record the request inputs (messages/prompt/embeddings input) on the span. + * Only called when `recordInputs` is enabled. + */ +export function addRequestAttributes( + span: Span, + inputs: unknown, + operationName: string, + enableTruncation: boolean, +): void { + if (!inputs || typeof inputs !== 'object') { + return; + } + const params = inputs as WorkersAiInput; + + // Store embeddings input on a separate attribute and do not truncate it + if (operationName === 'embeddings') { + const text = params.text; + + if (text == null) { + return; + } + if (typeof text === 'string' && text.length === 0) { + return; + } + if (Array.isArray(text) && text.length === 0) { + return; + } + + span.setAttribute(GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE, typeof text === 'string' ? text : JSON.stringify(text)); + return; + } + + const src = params.messages ?? params.prompt; + if (src == null) { + return; + } + if (Array.isArray(src) && src.length === 0) { + return; + } + + const { systemInstructions, filteredMessages } = extractSystemInstructions(src); + + if (systemInstructions) { + span.setAttribute(GEN_AI_SYSTEM_INSTRUCTIONS_ATTRIBUTE, systemInstructions); + } + + span.setAttribute( + GEN_AI_INPUT_MESSAGES_ATTRIBUTE, + enableTruncation ? getTruncatedJsonString(filteredMessages) : getJsonString(filteredMessages), + ); + + span.setAttribute( + GEN_AI_INPUT_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE, + Array.isArray(filteredMessages) ? filteredMessages.length : 1, + ); +} + +/** + * Record the response attributes (token usage, response text, tool calls) on the span. + */ +export function addResponseAttributes(span: Span, result: unknown, recordOutputs: boolean): void { + if (!result || typeof result !== 'object') { + return; + } + + // Raw `Response` objects (from `returnRawResponse`/`websocket`) cannot be introspected without consuming them. + if (typeof Response !== 'undefined' && result instanceof Response) { + return; + } + + const response = result as WorkersAiOutput; + + if (response.usage) { + setTokenUsageAttributes(span, response.usage.prompt_tokens, response.usage.completion_tokens); + } + + if (recordOutputs) { + if (typeof response.response === 'string') { + span.setAttribute(GEN_AI_RESPONSE_TEXT_ATTRIBUTE, response.response); + } else if (response.response != null) { + span.setAttribute(GEN_AI_RESPONSE_TEXT_ATTRIBUTE, JSON.stringify(response.response)); + } + + if (Array.isArray(response.tool_calls) && response.tool_calls.length > 0) { + span.setAttribute(GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE, JSON.stringify(response.tool_calls)); + } + } +} diff --git a/packages/core/test/lib/tracing/workers-ai-streaming.test.ts b/packages/core/test/lib/tracing/workers-ai-streaming.test.ts new file mode 100644 index 000000000000..16f7983b0d4e --- /dev/null +++ b/packages/core/test/lib/tracing/workers-ai-streaming.test.ts @@ -0,0 +1,114 @@ +import { describe, expect, it } from 'vitest'; +import type { Span } from '../../../src'; +import { + GEN_AI_RESPONSE_STREAMING_ATTRIBUTE, + GEN_AI_RESPONSE_TEXT_ATTRIBUTE, + GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE, + GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE, + GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE, +} from '../../../src/tracing/ai/gen-ai-attributes'; +import { instrumentWorkersAiStream } from '../../../src/tracing/workers-ai/streaming'; + +function createMockSpan(): { span: Span; attributes: Record; ended: () => boolean } { + const attributes: Record = {}; + let isEnded = false; + const span = { + isRecording: () => !isEnded, + setAttribute: (key: string, value: unknown) => { + attributes[key] = value; + }, + setAttributes: (attrs: Record) => { + Object.assign(attributes, attrs); + }, + setStatus: () => {}, + end: () => { + isEnded = true; + }, + } as unknown as Span; + return { span, attributes, ended: () => isEnded }; +} + +function streamFromChunks(chunks: string[]): ReadableStream { + const encoder = new TextEncoder(); + const queued = [...chunks]; + return new ReadableStream({ + pull(controller) { + const next = queued.shift(); + if (next === undefined) { + controller.close(); + return; + } + controller.enqueue(encoder.encode(next)); + }, + }); +} + +describe('instrumentWorkersAiStream', () => { + it('passes the original bytes through untouched', async () => { + const { span } = createMockSpan(); + const raw = 'data: {"response":"Hello"}\n\ndata: [DONE]\n\n'; + + const instrumented = instrumentWorkersAiStream(streamFromChunks([raw]), span, true); + const passedThrough = await new Response(instrumented).text(); + + expect(passedThrough).toBe(raw); + }); + + it('accumulates response text and usage across SSE events split across read boundaries', async () => { + const { span, attributes, ended } = createMockSpan(); + // The second event is deliberately split mid-line across two reads. + const chunks = [ + 'data: {"response":"The capital "}\n\ndata: {"respo', + 'nse":"of France "}\n\ndata: {"response":"is Paris."', + ',"usage":{"prompt_tokens":12,"completion_tokens":7,"total_tokens":19}}\n\ndata: [DONE]\n\n', + ]; + + const instrumented = instrumentWorkersAiStream(streamFromChunks(chunks), span, true); + await new Response(instrumented).text(); + + expect(attributes[GEN_AI_RESPONSE_STREAMING_ATTRIBUTE]).toBe(true); + expect(attributes[GEN_AI_RESPONSE_TEXT_ATTRIBUTE]).toBe('The capital of France is Paris.'); + expect(attributes[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE]).toBe(12); + expect(attributes[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE]).toBe(7); + expect(attributes[GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE]).toBe(19); + expect(ended()).toBe(true); + }); + + it('records usage but not response text when recordOutputs is false', async () => { + const { span, attributes } = createMockSpan(); + const chunks = [ + 'data: {"response":"secret"}\n\n', + 'data: {"response":"","usage":{"prompt_tokens":3,"completion_tokens":1,"total_tokens":4}}\n\ndata: [DONE]\n\n', + ]; + + const instrumented = instrumentWorkersAiStream(streamFromChunks(chunks), span, false); + await new Response(instrumented).text(); + + expect(attributes[GEN_AI_RESPONSE_TEXT_ATTRIBUTE]).toBeUndefined(); + expect(attributes[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE]).toBe(1); + }); + + it('ignores malformed SSE payloads without throwing', async () => { + const { span, attributes, ended } = createMockSpan(); + const chunks = ['data: not-json\n\n', 'data: {"response":"ok"}\n\ndata: [DONE]\n\n']; + + const instrumented = instrumentWorkersAiStream(streamFromChunks(chunks), span, true); + await new Response(instrumented).text(); + + expect(attributes[GEN_AI_RESPONSE_TEXT_ATTRIBUTE]).toBe('ok'); + expect(ended()).toBe(true); + }); + + it('ends the span when the consumer cancels the stream', async () => { + const { span, attributes, ended } = createMockSpan(); + + const instrumented = instrumentWorkersAiStream(streamFromChunks(['data: {"response":"partial"}\n\n']), span, true); + + const reader = instrumented.getReader(); + await reader.read(); + await reader.cancel('no longer needed'); + + expect(ended()).toBe(true); + expect(attributes[GEN_AI_RESPONSE_STREAMING_ATTRIBUTE]).toBe(true); + }); +}); diff --git a/packages/core/test/lib/tracing/workers-ai.test.ts b/packages/core/test/lib/tracing/workers-ai.test.ts new file mode 100644 index 000000000000..9c15e7c0b5a5 --- /dev/null +++ b/packages/core/test/lib/tracing/workers-ai.test.ts @@ -0,0 +1,41 @@ +import { describe, expect, it, vi } from 'vitest'; +import { instrumentWorkersAiClient } from '../../../src/tracing/workers-ai'; + +describe('instrumentWorkersAiClient', () => { + it('returns an already-instrumented client as-is instead of double-wrapping', () => { + const client = { run: vi.fn().mockResolvedValue({ response: 'ok' }) }; + + const instrumented = instrumentWorkersAiClient(client); + const instrumentedTwice = instrumentWorkersAiClient(instrumented); + + expect(instrumented).not.toBe(client); + expect(instrumentedTwice).toBe(instrumented); + }); + + it('passes through non-run methods bound to the original client', () => { + class FakeAi { + #internal = 'secret'; + + public run = vi.fn().mockResolvedValue({ response: 'ok' }); + + public readInternal(): string { + return this.#internal; + } + } + + const client = new FakeAi(); + const instrumented = instrumentWorkersAiClient(client); + + expect(instrumented.readInternal()).toBe('secret'); + }); + + it('creates a span and forwards arguments for run calls', async () => { + const client = { run: vi.fn().mockResolvedValue({ response: 'Paris' }) }; + const instrumented = instrumentWorkersAiClient(client); + + const result = await instrumented.run('@cf/meta/llama-3.1-8b-instruct', { prompt: 'Hello' }); + + expect(client.run).toHaveBeenCalledWith('@cf/meta/llama-3.1-8b-instruct', { prompt: 'Hello' }); + expect(result).toEqual({ response: 'Paris' }); + }); +}); diff --git a/packages/core/test/lib/utils/workers-ai-utils.test.ts b/packages/core/test/lib/utils/workers-ai-utils.test.ts new file mode 100644 index 000000000000..eed6cf170ed3 --- /dev/null +++ b/packages/core/test/lib/utils/workers-ai-utils.test.ts @@ -0,0 +1,242 @@ +import { describe, expect, it } from 'vitest'; +import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '../../../src'; +import type { Span } from '../../../src'; +import { + GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE, + GEN_AI_INPUT_MESSAGES_ATTRIBUTE, + GEN_AI_OPERATION_NAME_ATTRIBUTE, + GEN_AI_REQUEST_FREQUENCY_PENALTY_ATTRIBUTE, + GEN_AI_REQUEST_MAX_TOKENS_ATTRIBUTE, + GEN_AI_REQUEST_MODEL_ATTRIBUTE, + GEN_AI_REQUEST_PRESENCE_PENALTY_ATTRIBUTE, + GEN_AI_REQUEST_STREAM_ATTRIBUTE, + GEN_AI_REQUEST_TEMPERATURE_ATTRIBUTE, + GEN_AI_REQUEST_TOP_K_ATTRIBUTE, + GEN_AI_REQUEST_TOP_P_ATTRIBUTE, + GEN_AI_RESPONSE_TEXT_ATTRIBUTE, + GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE, + GEN_AI_SYSTEM_ATTRIBUTE, + GEN_AI_SYSTEM_INSTRUCTIONS_ATTRIBUTE, + GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE, + GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE, + GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE, +} from '../../../src/tracing/ai/gen-ai-attributes'; +import { WORKERS_AI_ORIGIN, WORKERS_AI_SYSTEM_NAME } from '../../../src/tracing/workers-ai/constants'; +import { + addRequestAttributes, + addResponseAttributes, + extractRequestAttributes, + getOperationName, +} from '../../../src/tracing/workers-ai/utils'; + +const MODEL = '@cf/meta/llama-3.1-8b-instruct'; + +function createMockSpan(): { span: Span; attributes: Record } { + const attributes: Record = {}; + const span = { + setAttribute: (key: string, value: unknown) => { + attributes[key] = value; + }, + setAttributes: (attrs: Record) => { + Object.assign(attributes, attrs); + }, + } as unknown as Span; + return { span, attributes }; +} + +describe('workers-ai utils', () => { + describe('getOperationName', () => { + it('returns "chat" for prompt inputs', () => { + expect(getOperationName({ prompt: 'Hello' })).toBe('chat'); + }); + + it('returns "chat" for messages inputs', () => { + expect(getOperationName({ messages: [{ role: 'user', content: 'Hi' }] })).toBe('chat'); + }); + + it('returns "embeddings" for text inputs', () => { + expect(getOperationName({ text: 'embed me' })).toBe('embeddings'); + }); + + it('prefers "chat" when both messages and text are present', () => { + expect(getOperationName({ messages: [{ role: 'user', content: 'Hi' }], text: 'embed me' })).toBe('chat'); + }); + + it('falls back to "chat" for null, undefined and empty inputs', () => { + expect(getOperationName(null)).toBe('chat'); + expect(getOperationName(undefined)).toBe('chat'); + expect(getOperationName({})).toBe('chat'); + }); + }); + + describe('extractRequestAttributes', () => { + it('sets exactly the base attributes for a minimal request', () => { + expect(extractRequestAttributes(MODEL, { prompt: 'Hello' }, 'chat')).toEqual({ + [GEN_AI_SYSTEM_ATTRIBUTE]: WORKERS_AI_SYSTEM_NAME, + [GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'chat', + [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: WORKERS_AI_ORIGIN, + [GEN_AI_REQUEST_MODEL_ATTRIBUTE]: MODEL, + }); + }); + + it('maps all supported request parameters', () => { + expect( + extractRequestAttributes( + MODEL, + { + prompt: 'Hello', + temperature: 0.5, + max_tokens: 100, + top_p: 0.9, + top_k: 40, + frequency_penalty: 0.1, + presence_penalty: 0.2, + stream: true, + }, + 'chat', + ), + ).toEqual({ + [GEN_AI_SYSTEM_ATTRIBUTE]: WORKERS_AI_SYSTEM_NAME, + [GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'chat', + [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: WORKERS_AI_ORIGIN, + [GEN_AI_REQUEST_MODEL_ATTRIBUTE]: MODEL, + [GEN_AI_REQUEST_TEMPERATURE_ATTRIBUTE]: 0.5, + [GEN_AI_REQUEST_MAX_TOKENS_ATTRIBUTE]: 100, + [GEN_AI_REQUEST_TOP_P_ATTRIBUTE]: 0.9, + [GEN_AI_REQUEST_TOP_K_ATTRIBUTE]: 40, + [GEN_AI_REQUEST_FREQUENCY_PENALTY_ATTRIBUTE]: 0.1, + [GEN_AI_REQUEST_PRESENCE_PENALTY_ATTRIBUTE]: 0.2, + [GEN_AI_REQUEST_STREAM_ATTRIBUTE]: true, + }); + }); + + it('does not set the stream attribute when stream is false', () => { + const attrs = extractRequestAttributes(MODEL, { prompt: 'Hello', stream: false }, 'chat'); + expect(attrs[GEN_AI_REQUEST_STREAM_ATTRIBUTE]).toBeUndefined(); + }); + + it('falls back to "unknown" model when model is not a string', () => { + const attrs = extractRequestAttributes(undefined, { prompt: 'Hello' }, 'chat'); + expect(attrs[GEN_AI_REQUEST_MODEL_ATTRIBUTE]).toBe('unknown'); + }); + }); + + describe('addRequestAttributes', () => { + it('records messages and extracts system instructions', () => { + const { span, attributes } = createMockSpan(); + + addRequestAttributes( + span, + { + messages: [ + { role: 'system', content: 'You are helpful.' }, + { role: 'user', content: 'Hi' }, + ], + }, + 'chat', + false, + ); + + expect(attributes[GEN_AI_SYSTEM_INSTRUCTIONS_ATTRIBUTE]).toBe( + JSON.stringify([{ type: 'text', content: 'You are helpful.' }]), + ); + expect(attributes[GEN_AI_INPUT_MESSAGES_ATTRIBUTE]).toBe(JSON.stringify([{ role: 'user', content: 'Hi' }])); + }); + + it('records the prompt string directly', () => { + const { span, attributes } = createMockSpan(); + + addRequestAttributes(span, { prompt: 'Hello world' }, 'chat', false); + + expect(attributes[GEN_AI_INPUT_MESSAGES_ATTRIBUTE]).toBe('Hello world'); + }); + + it('records embeddings input on a dedicated attribute', () => { + const { span, attributes } = createMockSpan(); + + addRequestAttributes(span, { text: ['embed a', 'embed b'] }, 'embeddings', false); + + expect(attributes).toEqual({ [GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE]: JSON.stringify(['embed a', 'embed b']) }); + }); + + it('records nothing for an empty messages array', () => { + const { span, attributes } = createMockSpan(); + + addRequestAttributes(span, { messages: [] }, 'chat', false); + + expect(attributes).toEqual({}); + }); + + it('records nothing for empty embeddings input', () => { + const { span, attributes } = createMockSpan(); + + addRequestAttributes(span, { text: '' }, 'embeddings', false); + + expect(attributes).toEqual({}); + }); + + it('records nothing when inputs are missing', () => { + const { span, attributes } = createMockSpan(); + + addRequestAttributes(span, undefined, 'chat', false); + + expect(attributes).toEqual({}); + }); + }); + + describe('addResponseAttributes', () => { + it('sets token usage and computes the total from input and output tokens', () => { + const { span, attributes } = createMockSpan(); + + addResponseAttributes(span, { response: 'Paris', usage: { prompt_tokens: 12, completion_tokens: 7 } }, false); + + expect(attributes).toEqual({ + [GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE]: 12, + [GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE]: 7, + [GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE]: 19, + }); + }); + + it('does not record response text when recordOutputs is false', () => { + const { span, attributes } = createMockSpan(); + + addResponseAttributes(span, { response: 'Paris' }, false); + + expect(attributes[GEN_AI_RESPONSE_TEXT_ATTRIBUTE]).toBeUndefined(); + }); + + it('records response text and tool calls when recordOutputs is true', () => { + const { span, attributes } = createMockSpan(); + const toolCalls = [{ name: 'lookup', arguments: { city: 'Paris' } }]; + + addResponseAttributes(span, { response: 'Paris', tool_calls: toolCalls }, true); + + expect(attributes[GEN_AI_RESPONSE_TEXT_ATTRIBUTE]).toBe('Paris'); + expect(attributes[GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE]).toBe(JSON.stringify(toolCalls)); + }); + + it('serializes non-string response payloads as JSON', () => { + const { span, attributes } = createMockSpan(); + + addResponseAttributes(span, { response: { translated_text: 'Bonjour' } }, true); + + expect(attributes[GEN_AI_RESPONSE_TEXT_ATTRIBUTE]).toBe(JSON.stringify({ translated_text: 'Bonjour' })); + }); + + it('ignores raw Response objects', () => { + const { span, attributes } = createMockSpan(); + + addResponseAttributes(span, new Response('raw'), true); + + expect(attributes).toEqual({}); + }); + + it('ignores non-object results', () => { + const { span, attributes } = createMockSpan(); + + addResponseAttributes(span, null, true); + + expect(attributes).toEqual({}); + }); + }); +});