diff --git a/packages/server-utils/src/ai/pi-ai/chat.ts b/packages/server-utils/src/ai/pi-ai/chat.ts new file mode 100644 index 000000000000..2e401efcba2e --- /dev/null +++ b/packages/server-utils/src/ai/pi-ai/chat.ts @@ -0,0 +1,123 @@ +import type { Span, StartSpanOptions } from '@sentry/core'; +import { + getCurrentScope, + getIsolationScope, + isURLObjectRelative, + parseStringToURLObject, + SPAN_STATUS_ERROR, + stringify, +} from '@sentry/core'; +import { + GEN_AI_CONVERSATION_ID, + GEN_AI_INPUT_MESSAGES, + GEN_AI_OPERATION_NAME, + GEN_AI_OUTPUT_MESSAGES, + GEN_AI_PROVIDER_NAME, + GEN_AI_REQUEST_MAX_TOKENS, + GEN_AI_REQUEST_MODEL, + GEN_AI_REQUEST_REASONING_LEVEL, + GEN_AI_REQUEST_TEMPERATURE, + GEN_AI_RESPONSE_FINISH_REASONS, + GEN_AI_RESPONSE_ID, + GEN_AI_RESPONSE_MODEL, + GEN_AI_RESPONSE_STREAMING, + GEN_AI_SYSTEM_INSTRUCTIONS, + GEN_AI_TOOL_DEFINITIONS, + SENTRY_ORIGIN, + SERVER_ADDRESS, + SERVER_PORT, +} from '@sentry/conventions/attributes'; +import { GEN_AI_CHAT } from '@sentry/conventions/op'; +import type { PiAiContext } from './messages'; +import { + piAiAssistantMessageToGenAiMessage, + piAiFinishReason, + piAiMessagesToGenAiMessages, + piAiSystemInstructions, + piAiToolDefinitions, +} from './messages'; +import type { PiAssistantMessage, PiModel, PiStreamOptions } from './types'; +import { setPiAiUsageAttributes } from './usage'; + +export interface PiAiChatRequest { + model: PiModel; + /** Absent for the deferred calls, which pass a handle in place of the request context. */ + context?: PiAiContext; + callOptions?: PiStreamOptions; + streaming: boolean; + origin: string; + recordInputs: boolean; +} + +/** + * The span options of the `gen_ai.chat` span of a pi-ai request. Returned rather than started so + * the caller picks whether the span is active while the request is sent. + */ +export function getPiAiChatSpanOptions(request: PiAiChatRequest): StartSpanOptions { + const { model, context, callOptions, streaming, origin, recordInputs } = request; + const conversationId = + getCurrentScope().getScopeData().conversationId ?? getIsolationScope().getScopeData().conversationId; + + return { + name: model.id ? `chat ${model.id}` : 'chat', + op: GEN_AI_CHAT, + attributes: { + [SENTRY_ORIGIN]: origin, + [GEN_AI_OPERATION_NAME]: 'chat', + [GEN_AI_PROVIDER_NAME]: model.provider, + [GEN_AI_REQUEST_MODEL]: model.id, + [GEN_AI_CONVERSATION_ID]: conversationId, + [GEN_AI_RESPONSE_STREAMING]: streaming || undefined, + [GEN_AI_REQUEST_TEMPERATURE]: callOptions?.temperature, + [GEN_AI_REQUEST_MAX_TOKENS]: callOptions?.maxTokens, + [GEN_AI_REQUEST_REASONING_LEVEL]: callOptions?.reasoning, + ...getServerAttributes(model), + ...(recordInputs && context ? getRequestContentAttributes(context) : {}), + }, + }; +} + +/** + * Record the settled message of a pi-ai request on its span. pi-ai never rejects a request: a + * provider failure or an abort settles as a message with `stopReason` `error` or `aborted`. + */ +export function setPiAiResponseAttributes(span: Span, message: PiAssistantMessage, recordOutputs: boolean): void { + const { responseModel, model, responseId, stopReason, usage } = message; + + const finishReason = piAiFinishReason(stopReason); + const output = recordOutputs ? piAiAssistantMessageToGenAiMessage(message, finishReason) : undefined; + + span.setAttributes({ + [GEN_AI_RESPONSE_MODEL]: responseModel ?? model, + [GEN_AI_RESPONSE_ID]: responseId, + [GEN_AI_RESPONSE_FINISH_REASONS]: finishReason ? stringify([finishReason]) : undefined, + [GEN_AI_OUTPUT_MESSAGES]: output ? stringify([output], String) : undefined, + }); + + const failed = stopReason === 'error'; + setPiAiUsageAttributes(span, usage, failed || stopReason === 'deferred'); + + if (failed) { + span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + } else if (stopReason === 'aborted') { + span.setStatus({ code: SPAN_STATUS_ERROR, message: 'cancelled' }); + } +} + +function getServerAttributes(model: PiModel): Record { + const url = model.baseUrl ? parseStringToURLObject(model.baseUrl) : undefined; + if (!url || isURLObjectRelative(url) || !url.hostname) { + return {}; + } + return { [SERVER_ADDRESS]: url.hostname, [SERVER_PORT]: url.port ? Number(url.port) : undefined }; +} + +function getRequestContentAttributes(context: PiAiContext): Record { + const messages = piAiMessagesToGenAiMessages(context.messages); + const tools = piAiToolDefinitions(context); + return { + [GEN_AI_SYSTEM_INSTRUCTIONS]: piAiSystemInstructions(context), + [GEN_AI_INPUT_MESSAGES]: messages.length ? stringify(messages) : undefined, + [GEN_AI_TOOL_DEFINITIONS]: tools.length ? stringify(tools) : undefined, + }; +} diff --git a/packages/server-utils/src/ai/pi-ai/types.ts b/packages/server-utils/src/ai/pi-ai/types.ts new file mode 100644 index 000000000000..f2355d1044d8 --- /dev/null +++ b/packages/server-utils/src/ai/pi-ai/types.ts @@ -0,0 +1,35 @@ +/** + * Structural subsets of the `@earendil-works/pi-ai` types the instrumentation reads. The SDK does + * not depend on the package, so only the fields used here are declared, and every one of them is + * treated as possibly absent at runtime. + */ + +import type { PiAiUsage } from './usage'; + +export interface PiAssistantMessage { + role?: string; + content?: unknown[]; + provider?: string; + model?: string; + responseModel?: string; + responseId?: string; + usage?: PiAiUsage; + stopReason?: string; + errorMessage?: string; +} + +export interface PiModel { + id?: string; + provider?: string; + baseUrl?: string; +} + +export interface PiStreamOptions { + temperature?: number; + maxTokens?: number; + reasoning?: string; +} + +export interface PiEventStream { + result(): Promise; +} diff --git a/packages/server-utils/src/ai/pi-durable/models.ts b/packages/server-utils/src/ai/pi-durable/models.ts index 196973159d70..07cf3a279165 100644 --- a/packages/server-utils/src/ai/pi-durable/models.ts +++ b/packages/server-utils/src/ai/pi-durable/models.ts @@ -1,48 +1,11 @@ -import type { Span } from '@sentry/core'; -import { - getCurrentScope, - getIsolationScope, - isThenable, - isURLObjectRelative, - parseStringToURLObject, - SPAN_STATUS_ERROR, - startSpanManual, - stringify, -} from '@sentry/core'; -import { - GEN_AI_CONVERSATION_ID, - GEN_AI_INPUT_MESSAGES, - GEN_AI_OPERATION_NAME, - GEN_AI_OUTPUT_MESSAGES, - GEN_AI_PROVIDER_NAME, - GEN_AI_REQUEST_MAX_TOKENS, - GEN_AI_REQUEST_MODEL, - GEN_AI_REQUEST_REASONING_LEVEL, - GEN_AI_REQUEST_TEMPERATURE, - GEN_AI_RESPONSE_FINISH_REASONS, - GEN_AI_RESPONSE_ID, - GEN_AI_RESPONSE_MODEL, - GEN_AI_RESPONSE_STREAMING, - GEN_AI_SYSTEM_INSTRUCTIONS, - GEN_AI_TOOL_DEFINITIONS, - SENTRY_ORIGIN, - SERVER_ADDRESS, - SERVER_PORT, -} from '@sentry/conventions/attributes'; -import { GEN_AI_CHAT } from '@sentry/conventions/op'; +import { isThenable, SPAN_STATUS_ERROR, startSpanManual } from '@sentry/core'; import type { GenAiOptions } from '../core/utils'; import { resolveAIRecordingOptions } from '../core/utils'; +import { getPiAiChatSpanOptions, setPiAiResponseAttributes } from '../pi-ai/chat'; import type { PiAiContext } from '../pi-ai/messages'; -import { - piAiAssistantMessageToGenAiMessage, - piAiFinishReason, - piAiMessagesToGenAiMessages, - piAiSystemInstructions, - piAiToolDefinitions, -} from '../pi-ai/messages'; -import { setPiAiUsageAttributes } from '../pi-ai/usage'; +import type { PiAssistantMessage, PiEventStream, PiModel, PiStreamOptions } from '../pi-ai/types'; import { PI_DURABLE_ORIGIN } from './constants'; -import type { PiAssistantMessage, PiEventStream, PiModel, PiModels, PiStreamOptions } from './types'; +import type { PiModels } from './types'; type ModelCall = (model: PiModel, request: unknown, options?: unknown) => unknown; @@ -107,27 +70,9 @@ function traceModelRequest( ): unknown { // Read per request: `dataCollection.genAI` belongs to the current client, which can change. const { recordInputs, recordOutputs } = resolveAIRecordingOptions(options); - const conversationId = - getCurrentScope().getScopeData().conversationId ?? getIsolationScope().getScopeData().conversationId; return startSpanManual( - { - name: model.id ? `chat ${model.id}` : 'chat', - op: GEN_AI_CHAT, - attributes: { - [SENTRY_ORIGIN]: PI_DURABLE_ORIGIN, - [GEN_AI_OPERATION_NAME]: 'chat', - [GEN_AI_PROVIDER_NAME]: model.provider, - [GEN_AI_REQUEST_MODEL]: model.id, - [GEN_AI_CONVERSATION_ID]: conversationId, - ...getServerAttributes(model), - ...(streaming ? { [GEN_AI_RESPONSE_STREAMING]: true } : {}), - [GEN_AI_REQUEST_TEMPERATURE]: callOptions?.temperature, - [GEN_AI_REQUEST_MAX_TOKENS]: callOptions?.maxTokens, - [GEN_AI_REQUEST_REASONING_LEVEL]: callOptions?.reasoning, - ...(recordInputs && context ? getRequestContentAttributes(context) : {}), - }, - }, + getPiAiChatSpanOptions({ model, context, callOptions, streaming, origin: PI_DURABLE_ORIGIN, recordInputs }), span => { let result: unknown; try { @@ -152,7 +97,10 @@ function traceModelRequest( return result; } (settled as Promise).then( - message => endChatSpan(span, message, recordOutputs), + message => { + setPiAiResponseAttributes(span, message, recordOutputs); + span.end(); + }, () => { span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); span.end(); @@ -163,51 +111,3 @@ function traceModelRequest( }, ); } - -function getServerAttributes(model: PiModel): Record { - const url = model.baseUrl ? parseStringToURLObject(model.baseUrl) : undefined; - if (!url || isURLObjectRelative(url) || !url.hostname) { - return {}; - } - return { [SERVER_ADDRESS]: url.hostname, ...(url.port ? { [SERVER_PORT]: Number(url.port) } : {}) }; -} - -function getRequestContentAttributes(context: PiAiContext): Record { - const messages = piAiMessagesToGenAiMessages(context.messages); - const tools = piAiToolDefinitions(context); - return { - [GEN_AI_SYSTEM_INSTRUCTIONS]: piAiSystemInstructions(context), - [GEN_AI_INPUT_MESSAGES]: messages.length ? stringify(messages) : undefined, - [GEN_AI_TOOL_DEFINITIONS]: tools.length ? stringify(tools) : undefined, - }; -} - -function endChatSpan(span: Span, message: PiAssistantMessage, recordOutputs: boolean): void { - const responseModel = message.responseModel ?? message.model; - if (responseModel) { - span.setAttribute(GEN_AI_RESPONSE_MODEL, responseModel); - } - if (message.responseId) { - span.setAttribute(GEN_AI_RESPONSE_ID, message.responseId); - } - - const finishReason = piAiFinishReason(message.stopReason); - if (finishReason) { - span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, stringify([finishReason])); - } - - const failed = message.stopReason === 'error'; - setPiAiUsageAttributes(span, message.usage, failed || message.stopReason === 'deferred'); - - const output = recordOutputs ? piAiAssistantMessageToGenAiMessage(message, finishReason) : undefined; - if (output) { - span.setAttribute(GEN_AI_OUTPUT_MESSAGES, stringify([output], String)); - } - - if (failed) { - span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); - } else if (message.stopReason === 'aborted') { - span.setStatus({ code: SPAN_STATUS_ERROR, message: 'cancelled' }); - } - span.end(); -} diff --git a/packages/server-utils/src/ai/pi-durable/types.ts b/packages/server-utils/src/ai/pi-durable/types.ts index c4ec68ceef8c..e4b79f7fcd42 100644 --- a/packages/server-utils/src/ai/pi-durable/types.ts +++ b/packages/server-utils/src/ai/pi-durable/types.ts @@ -5,35 +5,7 @@ */ import type { PiAiContext } from '../pi-ai/messages'; -import type { PiAiUsage } from '../pi-ai/usage'; - -export interface PiAssistantMessage { - role?: string; - content?: unknown[]; - provider?: string; - model?: string; - responseModel?: string; - responseId?: string; - usage?: PiAiUsage; - stopReason?: string; - errorMessage?: string; -} - -export interface PiModel { - id?: string; - provider?: string; - baseUrl?: string; -} - -export interface PiStreamOptions { - temperature?: number; - maxTokens?: number; - reasoning?: string; -} - -export interface PiEventStream { - result(): Promise; -} +import type { PiAssistantMessage, PiEventStream, PiModel, PiStreamOptions } from '../pi-ai/types'; /** * The `Models` methods that send a model request. Everything else passes through untouched. The diff --git a/packages/server-utils/test/ai/lib/tracing/pi-durable.test.ts b/packages/server-utils/test/ai/lib/tracing/pi-durable.test.ts index e7bd701485cf..ef8c7aeeff31 100644 --- a/packages/server-utils/test/ai/lib/tracing/pi-durable.test.ts +++ b/packages/server-utils/test/ai/lib/tracing/pi-durable.test.ts @@ -52,7 +52,6 @@ import { createRuns, startRun } from '../../../../src/ai/pi-durable/runs'; import { instrumentTool, markBuiltInTool } from '../../../../src/ai/pi-durable/tools'; import type { PiCommitChange, - PiEventStream, PiHarnessOptions, PiLiveState, PiModels, @@ -64,6 +63,7 @@ import type { PiToolExecutionApi, PiToolExecutionResult, } from '../../../../src/ai/pi-durable/types'; +import type { PiEventStream } from '../../../../src/ai/pi-ai/types'; import { ANTHROPIC_AI_INTEGRATION_NAME } from '../../../../src/ai/anthropic-ai/constants'; import { GOOGLE_GENAI_INTEGRATION_NAME } from '../../../../src/ai/google-genai/constants'; import { OPENAI_INTEGRATION_NAME } from '../../../../src/ai/openai/constants';