ai
Version:
AI SDK by Vercel - build apps like ChatGPT, Claude, Gemini, and more with a single interface for any model using the Vercel AI Gateway or go direct to OpenAI, Anthropic, Google, or any other model provider.
336 lines (310 loc) • 11.3 kB
text/typescript
import { asArray, type Context } from '@ai-sdk/provider-utils';
import type { Callback } from '../util/callback';
import { VERSION } from '../version';
import {
filterIncludedContext,
filterToolContext,
filterToolsContext,
} from './filter-included-context';
import { mergeCallbacks } from '../util/merge-callbacks';
import type {
InferTelemetryEvent,
Telemetry,
TelemetryDispatcher,
} from './telemetry';
import {
openTelemetryChannelSpanContext,
runWithTracingChannelSpan,
} from './tracing-channel-publisher';
import { getGlobalTelemetryIntegrations } from './telemetry-registry';
import type { TelemetryOptions } from './telemetry-options';
import { filterTelemetryEvent } from './filter-telemetry-event';
import { redactTelemetryCredentials } from './redact-telemetry-credentials';
/**
* The subset of `TelemetryDispatcher` keys whose values are Callback callbacks.
* This excludes dispatcher-only helpers and non-Callback properties such as
* `executeLanguageModelCall` and `executeTool`.
*/
type TelemetryCallbackKey = Exclude<
keyof {
[K in keyof TelemetryDispatcher as TelemetryDispatcher[K] extends
| Callback<any>
| undefined
? K
: never]: true;
},
'runInTracingChannelSpan' | 'startTracingChannelContext'
>;
/**
* Resolves the public event type accepted by a telemetry callback key.
*/
type TelemetryEvent<K extends TelemetryCallbackKey> =
TelemetryDispatcher[K] extends Callback<infer EVENT> | undefined
? EVENT
: never;
function augmentEvent<EVENT>(
event: EVENT,
telemetry: Pick<
TelemetryOptions,
| 'recordInputs'
| 'recordOutputs'
| 'functionId'
| 'includeRuntimeContext'
| 'includeToolsContext'
>,
filterContext = false,
): InferTelemetryEvent<EVENT> {
const augmentedEvent = Object.assign(
Object.create(Object.getPrototypeOf(event)),
event,
{
recordInputs: telemetry.recordInputs,
recordOutputs: telemetry.recordOutputs,
functionId: telemetry.functionId,
aiVersion: VERSION,
},
);
if (
filterContext &&
event != null &&
typeof event === 'object' &&
'runtimeContext' in event
) {
augmentedEvent.runtimeContext = filterIncludedContext({
context: event.runtimeContext as Context,
includeContext: telemetry.includeRuntimeContext,
});
}
if (filterContext && event != null && typeof event === 'object') {
if ('toolsContext' in event) {
augmentedEvent.toolsContext = filterToolsContext({
toolsContext: event.toolsContext as Record<string, Context>,
includeToolsContext: telemetry.includeToolsContext,
});
} else if (
'toolContext' in event &&
event.toolContext != null &&
'toolCall' in event &&
event.toolCall != null &&
typeof event.toolCall === 'object' &&
'toolName' in event.toolCall
) {
augmentedEvent.toolContext = filterToolContext({
toolName: event.toolCall.toolName as string,
toolContext: event.toolContext,
includeToolsContext: telemetry.includeToolsContext,
});
}
}
return redactTelemetryCredentials(
filterTelemetryEvent(augmentedEvent, telemetry),
);
}
/**
* Creates a telemetry dispatcher that sends telemetry events
* to the resolved set of integrations.
*
* When per-call integrations are provided, they take precedence over the globally
* registered integrations for that call. When no per-call integrations are
* provided, the globally registered integrations are used.
*
* @param args.telemetry - Optional per-call telemetry settings and integrations.
*
* @returns A telemetry dispatcher that fans out lifecycle events to the
* resolved set of integrations.
*/
export function createTelemetryDispatcher({
telemetry,
}: {
telemetry?: TelemetryOptions;
// operationId: string;
}): TelemetryDispatcher {
// When telemetry is explicitly disabled, return a dispatcher
// that performs no work and lets tool execution pass through unwrapped.
if (telemetry?.isEnabled === false) {
return {};
}
const localIntegrations = telemetry?.integrations;
const integrations: Array<Telemetry> =
localIntegrations != null
? asArray(localIntegrations)
: getGlobalTelemetryIntegrations();
const telemetryMetadata = {
recordInputs: telemetry?.recordInputs,
recordOutputs: telemetry?.recordOutputs,
functionId: telemetry?.functionId,
includeRuntimeContext: telemetry?.includeRuntimeContext,
includeToolsContext: telemetry?.includeToolsContext,
};
// Raw operation results may contain both inputs and outputs.
const recordResult =
telemetry?.recordInputs !== false && telemetry?.recordOutputs !== false;
const mergeTelemetryCallback = <KEY extends TelemetryCallbackKey>(
key: KEY,
deprecatedKey?: keyof Telemetry,
): Callback<TelemetryEvent<KEY>> | undefined => {
const integrationCallbacks = (
integrations
.map(integration => {
const callback =
integration[key] ??
(deprecatedKey == null ? undefined : integration[deprecatedKey]);
return typeof callback === 'function'
? callback.bind(integration)
: undefined;
})
.filter(Boolean) as Array<
Callback<InferTelemetryEvent<TelemetryEvent<KEY>>>
>
).map(
callback =>
((event: TelemetryEvent<KEY>) =>
callback(augmentEvent(event, telemetryMetadata))) as Callback<
TelemetryEvent<KEY>
>,
);
if (integrationCallbacks.length === 0) {
return undefined;
}
const mergedIntegrationCallback = mergeCallbacks(...integrationCallbacks);
return async (event: TelemetryEvent<KEY>) => {
await mergedIntegrationCallback(event);
};
};
const onStepEnd = mergeTelemetryCallback('onStepEnd');
const onStepFinish = mergeTelemetryCallback('onStepFinish');
const executeLanguageModelCallWrappers = integrations
.map(integration => integration.executeLanguageModelCall?.bind(integration))
.filter(Boolean) as Array<
NonNullable<Telemetry['executeLanguageModelCall']>
>;
const executeToolWrappers = integrations
.map(integration => integration.executeTool?.bind(integration))
.filter(Boolean) as Array<NonNullable<Telemetry['executeTool']>>;
return {
runInTracingChannelSpan: async ({ type, event, execute }) =>
await runWithTracingChannelSpan(
{
type,
event: augmentEvent(event, telemetryMetadata, true),
},
execute,
recordResult,
),
startTracingChannelContext: ({ type, event, completion }) =>
openTelemetryChannelSpanContext({
message: {
type,
event: augmentEvent(event, telemetryMetadata, true),
},
completion,
recordResult,
}),
onStart: mergeTelemetryCallback('onStart'),
onStepStart: mergeTelemetryCallback('onStepStart'),
onLanguageModelCallStart: mergeTelemetryCallback(
'onLanguageModelCallStart',
),
onLanguageModelCallEnd: mergeTelemetryCallback('onLanguageModelCallEnd'),
onToolExecutionStart: mergeTelemetryCallback('onToolExecutionStart'),
onToolExecutionEnd: mergeTelemetryCallback('onToolExecutionEnd'),
// Fan out step-end events to both the new `onStepEnd` callback and the
// deprecated `onStepFinish` callback so integrations that still implement
// only `onStepFinish` keep receiving step-end events during the deprecation
// window.
onStepEnd:
onStepEnd == null && onStepFinish == null
? undefined
: mergeCallbacks(onStepEnd, onStepFinish),
onStepError: mergeTelemetryCallback('onStepError'),
onObjectStepStart: mergeTelemetryCallback('onObjectStepStart'),
onObjectStepEnd: mergeTelemetryCallback('onObjectStepEnd'),
onEmbedStart: mergeTelemetryCallback('onEmbedStart'),
onEmbedEnd: mergeTelemetryCallback('onEmbedEnd'),
onRerankStart: mergeTelemetryCallback('onRerankStart'),
onRerankEnd: mergeTelemetryCallback('onRerankEnd'),
experimental_onDecideStart: mergeTelemetryCallback(
'experimental_onDecideStart',
'experimental_onEvaluateStart',
),
experimental_onDecisionModelCallStart: mergeTelemetryCallback(
'experimental_onDecisionModelCallStart',
'experimental_onEvaluationModelCallStart',
),
experimental_onDecisionModelCallEnd: mergeTelemetryCallback(
'experimental_onDecisionModelCallEnd',
'experimental_onEvaluationModelCallEnd',
),
experimental_onDecideEnd: mergeTelemetryCallback(
'experimental_onDecideEnd',
'experimental_onEvaluateEnd',
),
experimental_onStreamTranscriptionStart: mergeTelemetryCallback(
'experimental_onStreamTranscriptionStart',
),
experimental_onStreamTranscriptionEnd: mergeTelemetryCallback(
'experimental_onStreamTranscriptionEnd',
),
onEnd: mergeTelemetryCallback('onEnd'),
onAbort: mergeTelemetryCallback('onAbort'),
onError: mergeTelemetryCallback('onError'),
/**
* Runs provider calls inside integration-specific context so
* auto-instrumented provider requests can be associated with model work.
*/
executeLanguageModelCall: async ({ execute, ...event }) => {
const augmentedEvent = augmentEvent(event, telemetryMetadata);
let wrappedExecute = execute;
for (const executeWrapper of executeLanguageModelCallWrappers) {
const innerExecute = wrappedExecute;
wrappedExecute = async () => {
const originals = new Map<unknown, unknown>();
const redact = <T>(original: T): T => {
const redacted = redactTelemetryCredentials(original);
if (redacted !== original) originals.set(redacted, original);
return redacted;
};
const restore = <T>(value: T): T =>
(originals.has(value) ? originals.get(value) : value) as T;
try {
return restore(
await executeWrapper({
...augmentedEvent,
execute: async () => {
try {
return redact(await innerExecute());
} catch (error) {
throw redact(error);
}
},
}),
);
} catch (error) {
throw restore(error);
}
};
}
return await runWithTracingChannelSpan(
{ type: 'languageModelCall', event: augmentedEvent },
wrappedExecute,
recordResult,
);
},
/**
* Composes all `executeTool` wrappers around the original tool execution.
* Each wrapper receives an `execute` function that calls the next wrapper in
* the chain, so integrations can establish nested telemetry context before
* delegating to the underlying tool.
*/
executeTool: async ({ execute, ...event }) => {
const augmentedEvent = augmentEvent(event, telemetryMetadata);
let wrappedExecute = execute;
for (const executeWrapper of executeToolWrappers) {
const innerExecute = wrappedExecute;
wrappedExecute = () =>
executeWrapper({ ...augmentedEvent, execute: innerExecute });
}
return await wrappedExecute();
},
};
}