diff --git a/packages/constants/src/tools.ts b/packages/constants/src/tools.ts index 4f094c47cc..21e0db498d 100644 --- a/packages/constants/src/tools.ts +++ b/packages/constants/src/tools.ts @@ -7,7 +7,7 @@ import { MessageRole, ToolMessage, } from './legacyCompiler' -import { ToolCallContent as ToolRequest } from 'promptl-ai' +import type { ToolCallContent as ToolRequest } from 'promptl-ai' import { StreamType, ToolCallResponse, VercelProviderTool } from './index' import { ToolSource, ToolSourceData } from './toolSources' import type { Tool } from 'ai' diff --git a/packages/constants/src/tracing/index.ts b/packages/constants/src/tracing/index.ts index d774593ebf..8a39a9875b 100644 --- a/packages/constants/src/tracing/index.ts +++ b/packages/constants/src/tracing/index.ts @@ -41,6 +41,8 @@ export const ATTR_LATITUDE_INTERNAL = `${ATTR_LATITUDE}.internal` export const ATTR_LATITUDE_TYPE = `${ATTR_LATITUDE}.type` +export const ATTR_LATITUDE_PROMPT_PATH = `${ATTR_LATITUDE}.promptPath` + export const GEN_AI_TOOL_TYPE_VALUE_FUNCTION = 'function' export const ATTR_GEN_AI_TOOL_CALL_ARGUMENTS = 'gen_ai.tool.call.arguments' export const ATTR_GEN_AI_TOOL_RESULT_VALUE = 'gen_ai.tool.result.value' @@ -139,9 +141,13 @@ export const ATTR_LLM_RESPONSE_STOP_REASON = 'llm.response.stop_reason' export const ATTR_AI_OPERATION_ID = 'ai.operationId' export const AI_OPERATION_ID_VALUE_TOOL = 'ai.toolCall' export const AI_OPERATION_ID_VALUE_GENERATE_TEXT = 'ai.generateText' +export const AI_OPERATION_ID_VALUE_GENERATE_TEXT_DO_GENERATE = 'ai.generateText.doGenerate' // prettier-ignore export const AI_OPERATION_ID_VALUE_STREAM_TEXT = 'ai.streamText' +export const AI_OPERATION_ID_VALUE_STREAM_TEXT_DO_STREAM = 'ai.streamText.doStream' // prettier-ignore export const AI_OPERATION_ID_VALUE_GENERATE_OBJECT = 'ai.generateObject' +export const AI_OPERATION_ID_VALUE_GENERATE_OBJECT_DO_GENERATE = 'ai.generateObject.doGenerate' // prettier-ignore export const AI_OPERATION_ID_VALUE_STREAM_OBJECT = 'ai.streamObject' +export const AI_OPERATION_ID_VALUE_STREAM_OBJECT_DO_STREAM = 'ai.streamObject.doStream' // prettier-ignore export const ATTR_AI_TOOL_CALL_NAME = 'ai.toolCall.name' export const ATTR_AI_TOOL_CALL_ID = 'ai.toolCall.id' diff --git a/packages/constants/src/tracing/span.ts b/packages/constants/src/tracing/span.ts index 6a811bd8ca..ecb92797f9 100644 --- a/packages/constants/src/tracing/span.ts +++ b/packages/constants/src/tracing/span.ts @@ -1,5 +1,5 @@ import { FinishReason } from 'ai' -import { Message } from 'promptl-ai' +import type { Message } from 'promptl-ai' import { LogSources } from '../models' export enum SpanKind { @@ -134,14 +134,15 @@ export type ToolSpanMetadata = BaseSpanMetadata & { } export type PromptSpanMetadata = BaseSpanMetadata & { - experimentUuid: string - externalId: string + experimentUuid?: string + externalId?: string parameters: Record - promptUuid: string - template: string + promptUuid?: string // Document UUID (may be resolved from promptPath server-side) + promptPath?: string // Path-based prompt identification + template?: string versionUuid: string - source: LogSources - projectId: number + source?: LogSources + projectId?: number } export type CompletionSpanMetadata = BaseSpanMetadata & { diff --git a/packages/core/src/events/handlers/evaluateLiveLog.ts b/packages/core/src/events/handlers/evaluateLiveLog.ts index ba2d50c620..6d30f81f8c 100644 --- a/packages/core/src/events/handlers/evaluateLiveLog.ts +++ b/packages/core/src/events/handlers/evaluateLiveLog.ts @@ -50,6 +50,11 @@ export const evaluateLiveLogJob = async ({ return } + // Cannot evaluate logs without a document UUID + if (!promptSpanMetadata.promptUuid) { + return + } + const commitsRepository = new CommitsRepository(workspace.id) const commit = await commitsRepository .getCommitByUuid({ uuid: promptSpanMetadata.versionUuid }) diff --git a/packages/core/src/events/handlers/requestDocumentSuggestionJob.ts b/packages/core/src/events/handlers/requestDocumentSuggestionJob.ts index 7fde1fd688..3cc3bbce9b 100644 --- a/packages/core/src/events/handlers/requestDocumentSuggestionJob.ts +++ b/packages/core/src/events/handlers/requestDocumentSuggestionJob.ts @@ -29,7 +29,8 @@ export const requestDocumentSuggestionJobV2 = async ({ if (!workspace) throw new NotFoundError(`Workspace not found ${workspaceId}`) if (result.hasPassed || result.error || result.usedForSuggestion) return if (!evaluation.enableSuggestions) return - if (!LIVE_SUGGESTION_SOURCES.includes((metadata as PromptSpanMetadata).source)) return // prettier-ignore + const promptMetadata = metadata as PromptSpanMetadata + if (!promptMetadata.source || !LIVE_SUGGESTION_SOURCES.includes(promptMetadata.source)) return // prettier-ignore const { documentSuggestionsQueue } = await queues() documentSuggestionsQueue.add( diff --git a/packages/core/src/services/tracing/spans/completion.ts b/packages/core/src/services/tracing/spans/completion.ts index bb2e13368e..99483a3193 100644 --- a/packages/core/src/services/tracing/spans/completion.ts +++ b/packages/core/src/services/tracing/spans/completion.ts @@ -237,6 +237,22 @@ function extractConfiguration( return Result.ok(toCamelCase(configuration)) } +/** + * Safely parse JSON payloads which can come either as JSON string or an object. + */ +function parseJsonPayload(value: unknown): Record { + if (value === undefined || value === null) return {} + if (typeof value === 'object') return value as Record + if (typeof value === 'string') { + try { + return JSON.parse(value) + } catch { + return {} + } + } + return {} +} + function convertToolCalls( raws: Record[], ): TypedResult { @@ -252,14 +268,18 @@ function convertToolCalls( type: ContentType.toolCall, toolCallId: String(toolCall.id || ''), toolName: String(func.name || ''), - toolArguments: JSON.parse(String(func.arguments || '{}')), + toolArguments: parseJsonPayload(func.arguments), }) } else { + // Handles multiple formats: + // - OpenAI: { id, name, arguments (string) } + // - Anthropic: { toolUseId, toolName, input (object) } + // - Vercel AI SDK: { toolCallId, toolName, input (object) } toolCalls.push({ type: ContentType.toolCall, toolCallId: String(toolCall.id || toolCall.toolCallId || toolCall.toolUseId || ''), // prettier-ignore toolName: String(toolCall.name || toolCall.toolName || ''), - toolArguments: JSON.parse(String(toolCall.arguments || toolCall.toolArguments || toolCall.input || '{}')), // prettier-ignore + toolArguments: parseJsonPayload(toolCall.arguments || toolCall.toolArguments || toolCall.input), // prettier-ignore }) } } diff --git a/packages/core/src/services/tracing/spans/process.ts b/packages/core/src/services/tracing/spans/process.ts index 9e3f49d2a0..5ca198c056 100644 --- a/packages/core/src/services/tracing/spans/process.ts +++ b/packages/core/src/services/tracing/spans/process.ts @@ -9,9 +9,13 @@ import { import { database } from '../../../client' import { AI_OPERATION_ID_VALUE_GENERATE_OBJECT, + AI_OPERATION_ID_VALUE_GENERATE_OBJECT_DO_GENERATE, AI_OPERATION_ID_VALUE_GENERATE_TEXT, + AI_OPERATION_ID_VALUE_GENERATE_TEXT_DO_GENERATE, AI_OPERATION_ID_VALUE_STREAM_OBJECT, + AI_OPERATION_ID_VALUE_STREAM_OBJECT_DO_STREAM, AI_OPERATION_ID_VALUE_STREAM_TEXT, + AI_OPERATION_ID_VALUE_STREAM_TEXT_DO_STREAM, AI_OPERATION_ID_VALUE_TOOL, ATTR_AI_OPERATION_ID, ATTR_LATITUDE_INTERNAL, @@ -142,6 +146,10 @@ export function extractSpanType( case AI_OPERATION_ID_VALUE_STREAM_TEXT: case AI_OPERATION_ID_VALUE_GENERATE_OBJECT: case AI_OPERATION_ID_VALUE_STREAM_OBJECT: + case AI_OPERATION_ID_VALUE_GENERATE_TEXT_DO_GENERATE: + case AI_OPERATION_ID_VALUE_STREAM_TEXT_DO_STREAM: + case AI_OPERATION_ID_VALUE_GENERATE_OBJECT_DO_GENERATE: + case AI_OPERATION_ID_VALUE_STREAM_OBJECT_DO_STREAM: return Result.ok(SpanType.Completion) } diff --git a/packages/core/src/services/tracing/spans/processBulk.ts b/packages/core/src/services/tracing/spans/processBulk.ts index 0303e75113..5dbff81db3 100644 --- a/packages/core/src/services/tracing/spans/processBulk.ts +++ b/packages/core/src/services/tracing/spans/processBulk.ts @@ -16,9 +16,13 @@ import { cache as redis } from '../../../cache' import { database } from '../../../client' import { AI_OPERATION_ID_VALUE_GENERATE_OBJECT, + AI_OPERATION_ID_VALUE_GENERATE_OBJECT_DO_GENERATE, AI_OPERATION_ID_VALUE_GENERATE_TEXT, + AI_OPERATION_ID_VALUE_GENERATE_TEXT_DO_GENERATE, AI_OPERATION_ID_VALUE_STREAM_OBJECT, + AI_OPERATION_ID_VALUE_STREAM_OBJECT_DO_STREAM, AI_OPERATION_ID_VALUE_STREAM_TEXT, + AI_OPERATION_ID_VALUE_STREAM_TEXT_DO_STREAM, AI_OPERATION_ID_VALUE_TOOL, ATTR_AI_OPERATION_ID, ATTR_LATITUDE_TYPE, @@ -579,6 +583,10 @@ export function extractSpanType( case AI_OPERATION_ID_VALUE_STREAM_TEXT: case AI_OPERATION_ID_VALUE_GENERATE_OBJECT: case AI_OPERATION_ID_VALUE_STREAM_OBJECT: + case AI_OPERATION_ID_VALUE_GENERATE_TEXT_DO_GENERATE: + case AI_OPERATION_ID_VALUE_STREAM_TEXT_DO_STREAM: + case AI_OPERATION_ID_VALUE_GENERATE_OBJECT_DO_GENERATE: + case AI_OPERATION_ID_VALUE_STREAM_OBJECT_DO_STREAM: return Result.ok(SpanType.Completion) } diff --git a/packages/core/src/services/tracing/spans/prompt.ts b/packages/core/src/services/tracing/spans/prompt.ts index 853160cf9c..c906d702d7 100644 --- a/packages/core/src/services/tracing/spans/prompt.ts +++ b/packages/core/src/services/tracing/spans/prompt.ts @@ -2,11 +2,17 @@ import { database } from '../../../client' import { ATTR_GEN_AI_REQUEST_PARAMETERS, ATTR_GEN_AI_REQUEST_TEMPLATE, + ATTR_LATITUDE_PROMPT_PATH, + HEAD_COMMIT, LogSources, SPAN_SPECIFICATIONS, SpanType, } from '../../../constants' import { Result } from '../../../lib/Result' +import { + CommitsRepository, + DocumentVersionsRepository, +} from '../../../repositories' import { SpanProcessArgs } from './shared' const specification = SPAN_SPECIFICATIONS[SpanType.Prompt] @@ -16,8 +22,8 @@ export const PromptSpanSpecification = { } async function process( - { attributes }: SpanProcessArgs, - _ = database, + { attributes, workspace }: SpanProcessArgs, + db = database, ) { let parameters: Record try { @@ -28,17 +34,84 @@ async function process( parameters = {} } - return Result.ok({ + // Get promptUuid from attributes, or try to resolve from promptPath + let promptUuid = attributes['latitude.documentUuid'] as string | undefined + const promptPath = attributes[ATTR_LATITUDE_PROMPT_PATH] as string | undefined + const projectId = attributes['latitude.projectId'] as number | undefined + const versionUuid = + (attributes['latitude.commitUuid'] as string) || HEAD_COMMIT + + // If promptPath is provided but promptUuid is not, resolve it + if (promptPath && !promptUuid && projectId) { + const resolvedUuid = await resolvePromptPathToUuid({ + promptPath, + projectId, + versionUuid, + workspaceId: workspace.id, + db, + }) + if (resolvedUuid) promptUuid = resolvedUuid + } + + const result = { parameters, template: attributes[ATTR_GEN_AI_REQUEST_TEMPLATE] as string, externalId: attributes['latitude.externalId'] as string, // References experimentUuid: attributes['latitude.experimentUuid'] as string, - promptUuid: attributes['latitude.documentUuid'] as string, - versionUuid: attributes['latitude.commitUuid'] as string, + promptUuid, + promptPath, + versionUuid, documentLogUuid: attributes['latitude.documentLogUuid'] as string, - projectId: attributes['latitude.projectId'] as number, + projectId, source: attributes['latitude.source'] as LogSources, - }) + } + + return Result.ok(result) +} + +/** + * Resolves a prompt path to its document UUID by looking up the document + * in the specified project and version. + */ +async function resolvePromptPathToUuid({ + promptPath, + projectId, + versionUuid, + workspaceId, + db, +}: { + promptPath: string + projectId: number + versionUuid: string + workspaceId: number + db: typeof database +}): Promise { + try { + // Get the commit + const commitsRepo = new CommitsRepository(workspaceId, db) + const commitResult = await commitsRepo.getCommitByUuid({ + uuid: versionUuid, + projectId, + }) + if (commitResult.error) { + return undefined + } + const commit = commitResult.value + + // Get the document by path + const docsRepo = new DocumentVersionsRepository(workspaceId, db) + const docResult = await docsRepo.getDocumentByPath({ + commit, + path: promptPath, + }) + if (docResult.error) { + return undefined + } + + return docResult.value.documentUuid + } catch { + return undefined + } } diff --git a/packages/telemetry/typescript/src/index.ts b/packages/telemetry/typescript/src/index.ts index 7db67b18a1..cdea6e6828 100644 --- a/packages/telemetry/typescript/src/index.ts +++ b/packages/telemetry/typescript/src/index.ts @@ -1 +1,2 @@ +export * from './processors' export * from './sdk' diff --git a/packages/telemetry/typescript/src/instrumentations/manual.ts b/packages/telemetry/typescript/src/instrumentations/manual.ts index bbe6ea8606..52d4ce40e1 100644 --- a/packages/telemetry/typescript/src/instrumentations/manual.ts +++ b/packages/telemetry/typescript/src/instrumentations/manual.ts @@ -30,6 +30,7 @@ import { ATTR_HTTP_REQUEST_URL, ATTR_HTTP_RESPONSE_BODY, ATTR_HTTP_RESPONSE_HEADER, + ATTR_LATITUDE_PROMPT_PATH, ATTR_LATITUDE_TYPE, GEN_AI_TOOL_TYPE_VALUE_FUNCTION, HEAD_COMMIT, @@ -125,11 +126,12 @@ export type EndHttpSpanOptions = EndSpanOptions & { export type PromptSpanOptions = StartSpanOptions & { documentLogUuid?: string // TODO(tracing): temporal related log, remove when observability is ready versionUuid?: string // Alias for commitUuid - promptUuid: string // Alias for documentUuid + promptUuid?: string // Alias for documentUuid + promptPath?: string // Path-based prompt identification (resolved server-side) projectId?: string experimentUuid?: string externalId?: string - template: string + template?: string parameters?: Record source?: LogSources } @@ -721,6 +723,7 @@ export class ManualInstrumentation implements BaseInstrumentation { documentLogUuid, versionUuid, promptUuid, + promptPath, projectId, experimentUuid, externalId, @@ -739,11 +742,12 @@ export class ManualInstrumentation implements BaseInstrumentation { } const attributes = { - [ATTR_GEN_AI_REQUEST_TEMPLATE]: template, + ...(template && { [ATTR_GEN_AI_REQUEST_TEMPLATE]: template }), [ATTR_GEN_AI_REQUEST_PARAMETERS]: jsonParameters, ['latitude.commitUuid']: versionUuid || HEAD_COMMIT, - ['latitude.documentUuid']: promptUuid, - ['latitude.projectId']: projectId, + ...(promptUuid && { ['latitude.documentUuid']: promptUuid }), + ...(promptPath && { [ATTR_LATITUDE_PROMPT_PATH]: promptPath }), + ...(projectId && { ['latitude.projectId']: projectId }), ...(documentLogUuid && { ['latitude.documentLogUuid']: documentLogUuid }), ...(experimentUuid && { ['latitude.experimentUuid']: experimentUuid }), ...(externalId && { ['latitude.externalId']: externalId }), @@ -751,7 +755,13 @@ export class ManualInstrumentation implements BaseInstrumentation { ...(rest.attributes || {}), } - return this.span(ctx, name || `prompt-${promptUuid}`, SpanType.Prompt, { + const spanName = + name || + (promptPath + ? `prompt-${promptPath}` + : `prompt-${promptUuid || 'external'}`) + + return this.span(ctx, spanName, SpanType.Prompt, { attributes, }) } diff --git a/packages/telemetry/typescript/src/processors/index.ts b/packages/telemetry/typescript/src/processors/index.ts new file mode 100644 index 0000000000..da9e476e42 --- /dev/null +++ b/packages/telemetry/typescript/src/processors/index.ts @@ -0,0 +1 @@ +export { NormalizingSpanProcessor } from './normalize' diff --git a/packages/telemetry/typescript/src/processors/normalize.ts b/packages/telemetry/typescript/src/processors/normalize.ts new file mode 100644 index 0000000000..ce089cdc5a --- /dev/null +++ b/packages/telemetry/typescript/src/processors/normalize.ts @@ -0,0 +1,85 @@ +import * as otel from '@opentelemetry/api' +import { + ReadableSpan, + Span, + SpanProcessor, +} from '@opentelemetry/sdk-trace-node' + +/** + * Maps Vercel AI SDK operation IDs to standard GenAI operation names + */ +const OPERATION_ID_MAPPINGS: Record = { + 'ai.generateText': 'chat', + 'ai.generateText.doGenerate': 'chat', + 'ai.streamText': 'chat', + 'ai.streamText.doStream': 'chat', + 'ai.generateObject': 'chat', + 'ai.generateObject.doGenerate': 'chat', + 'ai.streamObject': 'chat', + 'ai.streamObject.doStream': 'chat', + 'ai.embed': 'embeddings', + 'ai.embedMany': 'embeddings', + 'ai.toolCall': 'execute_tool', +} + +/** + * Maps operation IDs to Latitude span types + */ +const OPERATION_TO_LATITUDE_TYPE: Record = { + 'ai.generateText': 'completion', + 'ai.generateText.doGenerate': 'completion', + 'ai.streamText': 'completion', + 'ai.streamText.doStream': 'completion', + 'ai.generateObject': 'completion', + 'ai.generateObject.doGenerate': 'completion', + 'ai.streamObject': 'completion', + 'ai.streamObject.doStream': 'completion', + 'ai.embed': 'embedding', + 'ai.embedMany': 'embedding', + 'ai.toolCall': 'tool', +} + +/** + * A SpanProcessor that normalizes Vercel AI SDK telemetry to standard + * OpenTelemetry GenAI semantic conventions. + * + * This ensures a consistent experience regardless of which SDK generated the spans. + */ +export class NormalizingSpanProcessor implements SpanProcessor { + onStart(span: Span, _parentContext: otel.Context): void { + // We normalize on start so baggage and other processors can work with normalized attributes + const operationId = span.attributes['ai.operationId'] as string | undefined + + if (!operationId) { + // Not a Vercel AI SDK span, skip normalization + return + } + + // Map the operation ID to standard gen_ai.operation.name + const normalizedOperation = OPERATION_ID_MAPPINGS[operationId] + if (normalizedOperation) { + span.setAttribute('gen_ai.operation.name', normalizedOperation) + } + + // Set latitude.type if not already set + if (!span.attributes['latitude.type']) { + const latitudeType = OPERATION_TO_LATITUDE_TYPE[operationId] + if (latitudeType) { + span.setAttribute('latitude.type', latitudeType) + } + } + } + + onEnd(_span: ReadableSpan): void { + // Note: ReadableSpan attributes are immutable, so we do normalization on start + // If we need to normalize based on final attributes, we'd need a different approach + } + + forceFlush(): Promise { + return Promise.resolve() + } + + shutdown(): Promise { + return Promise.resolve() + } +} diff --git a/packages/telemetry/typescript/src/sdk/sdk.ts b/packages/telemetry/typescript/src/sdk/sdk.ts index d9c0eddf9d..7f8c1bb84e 100644 --- a/packages/telemetry/typescript/src/sdk/sdk.ts +++ b/packages/telemetry/typescript/src/sdk/sdk.ts @@ -10,8 +10,11 @@ import { StartSpanOptions, StartToolSpanOptions, } from '$telemetry/instrumentations' +import { NormalizingSpanProcessor } from '$telemetry/processors' import { DEFAULT_REDACT_SPAN_PROCESSOR } from '$telemetry/sdk/redact' import { + ATTR_LATITUDE_PROMPT_PATH, + HEAD_COMMIT, InstrumentationScope, SCOPE_LATITUDE, TraceContext, @@ -142,6 +145,30 @@ export type TelemetryOptions = { exporter?: SpanExporter processors?: SpanProcessor[] propagators?: TextMapPropagator[] + /** Enable debug logging to see all spans being created */ + debug?: boolean +} + +/** + * Options for the trace() method that wraps user code in a scoped context. + * All child spans created within the trace callback will inherit these metadata + * via OpenTelemetry baggage propagation. + */ +export type TraceOptions = { + /** Optional name for the trace span */ + name?: string + /** Project ID for the trace */ + projectId?: number | string + /** Version UUID (commit UUID) for the trace */ + versionUuid?: string + /** Path-based prompt identification (resolved server-side to documentUuid) */ + promptPath?: string + /** UUID-based prompt identification */ + promptUuid?: string + /** External identifier for correlation */ + externalId?: string + /** Additional custom metadata */ + metadata?: Record } export class LatitudeTelemetry { @@ -180,6 +207,9 @@ export class LatitudeTelemetry { new BaggageSpanProcessor(ALLOW_ALL_BAGGAGE_KEYS), ) + // Normalize Vercel AI SDK spans to standard GenAI semantic conventions + this.provider.addSpanProcessor(new NormalizingSpanProcessor()) + if (this.options.processors) { this.options.processors.forEach((processor) => { this.provider.addSpanProcessor(processor) @@ -430,6 +460,107 @@ export class LatitudeTelemetry { step(ctx: otel.Context, options?: StartSpanOptions) { return this.telemetry.step(ctx, options) } + + /** + * Wraps a function execution in a trace scope. + * All OpenTelemetry spans created within the callback will be children + * of this trace and inherit the metadata via baggage propagation. + * + * @example + * ```typescript + * const result = await telemetry.trace({ + * projectId: 123, + * versionUuid: 'abc-123', + * promptPath: 'chat/greeting', + * }, async () => { + * // All spans created here (e.g., OpenAI calls) will be children of this trace + * return openai.chat.completions.create({ ... }) + * }) + * ``` + */ + async trace(options: TraceOptions, fn: () => Promise): Promise { + const { + name, + projectId, + versionUuid, + promptPath, + promptUuid, + externalId, + metadata, + } = options + + // Create parent "prompt" span with latitude.* attributes + const span = this.telemetry.prompt(context.active(), { + name, + projectId: projectId?.toString(), + versionUuid: versionUuid || HEAD_COMMIT, + promptUuid, + promptPath, + externalId, + attributes: metadata as otel.Attributes, + }) + + // Set baggage for metadata propagation to child spans + // The BaggageSpanProcessor will copy these to span attributes + let ctx = span.context + const baggageEntries: Record = {} + + if (projectId) { + baggageEntries['latitude.projectId'] = { value: String(projectId) } + } + if (versionUuid) { + baggageEntries['latitude.commitUuid'] = { value: versionUuid } + } else { + baggageEntries['latitude.commitUuid'] = { value: HEAD_COMMIT } + } + if (promptUuid) { + baggageEntries['latitude.documentUuid'] = { value: promptUuid } + } + if (promptPath) { + baggageEntries[ATTR_LATITUDE_PROMPT_PATH] = { value: promptPath } + } + if (externalId) { + baggageEntries['latitude.externalId'] = { value: externalId } + } + + const baggage = propagation.createBaggage(baggageEntries) + ctx = propagation.setBaggage(ctx, baggage) + + try { + // Execute the function within this context + // All child spans will be parented to our span and inherit baggage + const result = await context.with(ctx, fn) + span.end() + return result + } catch (error) { + span.fail(error as Error) + throw error + } + } + + /** + * Creates a higher-order function wrapper for tracing. + * Returns a new function that, when called, automatically traces its execution + * with the provided options. + * + * @example + * ```typescript + * const tracedGenerate = telemetry.wrap(generateAIResponse, { + * projectId: 123, + * versionUuid: 'abc-123', + * promptPath: 'chat/greeting', + * }) + * + * // Later, each call is automatically traced: + * const result = await tracedGenerate(prompt, options) + * ``` + */ + wrap( + fn: (...args: TArgs) => Promise, + options: TraceOptions, + ): (...args: TArgs) => Promise { + return (...args: TArgs) => this.trace(options, () => fn(...args)) + } } export type { diff --git a/packages/telemetry/typescript/tests/telemetry/trace.test.ts b/packages/telemetry/typescript/tests/telemetry/trace.test.ts new file mode 100644 index 0000000000..027ba35087 --- /dev/null +++ b/packages/telemetry/typescript/tests/telemetry/trace.test.ts @@ -0,0 +1,468 @@ +import { LatitudeTelemetry } from '$telemetry/index' +import { ATTR_LATITUDE_PROMPT_PATH, HEAD_COMMIT } from '@latitude-data/constants' +import { context } from '@opentelemetry/api' +import { ReadableSpan } from '@opentelemetry/sdk-trace-node' +import { setupServer } from 'msw/node' +import { + afterAll, + afterEach, + beforeAll, + describe, + expect, + it, + vi, +} from 'vitest' +import { mockRequest, MockSpanProcessor } from '../utils' + +vi.hoisted(() => { + process.env.GATEWAY_BASE_URL = 'https://fake-host.com' + process.env.npm_package_name = 'fake-service-name' + process.env.npm_package_version = 'fake-scope-version' +}) + +describe('telemetry.trace', () => { + const gatewayMock = setupServer() + + beforeAll(() => { + gatewayMock.listen() + }) + + afterEach(() => { + gatewayMock.resetHandlers() + vi.clearAllMocks() + }) + + afterAll(() => { + gatewayMock.close() + }) + + it( + 'creates parent span with correct latitude attributes', + gatewayMock.boundary(async () => { + const { bodyMock } = mockRequest({ + server: gatewayMock, + method: 'post', + endpoint: '/api/v3/traces', + }) + + const sdk = new LatitudeTelemetry('fake-api-key', { + disableBatch: true, + }) + + await sdk.trace( + { + name: 'test-trace', + projectId: 123, + versionUuid: 'version-uuid-123', + promptPath: 'chat/greeting', + externalId: 'external-123', + }, + async () => { + // Simulate some work + await new Promise((resolve) => setTimeout(resolve, 10)) + return 'result' + }, + ) + + await sdk.shutdown() + + expect(bodyMock).toHaveBeenCalled() + const body = bodyMock.mock.calls[0]![0] + + // Find the prompt span (parent) + const spans = body.resourceSpans[0].scopeSpans[0].spans + expect(spans).toHaveLength(1) + + const promptSpan = spans[0] + expect(promptSpan.name).toBe('test-trace') + + // Verify latitude attributes + const attributes = promptSpan.attributes + expect(attributes).toContainEqual({ + key: 'latitude.type', + value: { stringValue: 'prompt' }, + }) + expect(attributes).toContainEqual({ + key: 'latitude.projectId', + value: { stringValue: '123' }, + }) + expect(attributes).toContainEqual({ + key: 'latitude.commitUuid', + value: { stringValue: 'version-uuid-123' }, + }) + expect(attributes).toContainEqual({ + key: ATTR_LATITUDE_PROMPT_PATH, + value: { stringValue: 'chat/greeting' }, + }) + expect(attributes).toContainEqual({ + key: 'latitude.externalId', + value: { stringValue: 'external-123' }, + }) + }), + ) + + it( + 'child spans inherit parent and baggage attributes', + gatewayMock.boundary(async () => { + const processorMock = new MockSpanProcessor() + + const sdk = new LatitudeTelemetry('fake-api-key', { + disableBatch: true, + processors: [processorMock], + }) + + mockRequest({ + server: gatewayMock, + method: 'post', + endpoint: '/api/v3/traces', + }) + + await sdk.trace( + { + name: 'parent-trace', + projectId: 456, + versionUuid: 'parent-version', + promptPath: 'test/path', + }, + async () => { + // Create a child completion span inside the trace + const completion = sdk.completion(context.active(), { + provider: 'openai', + model: 'gpt-4o', + configuration: { model: 'gpt-4o' }, + input: [{ role: 'user', content: 'Hello' }], + }) + completion.end({ + output: [{ role: 'assistant', content: 'Hi' }], + tokens: { prompt: 5, cached: 0, reasoning: 0, completion: 2 }, + finishReason: 'stop', + }) + return 'done' + }, + ) + + await sdk.shutdown() + + // Get all spans that were created + const endedSpans = processorMock.onEnd.mock.calls.map( + (call) => call[0] as ReadableSpan, + ) + expect(endedSpans.length).toBe(2) + + // Find parent and child spans + const parentSpan = endedSpans.find((s) => s.name === 'parent-trace') + const childSpan = endedSpans.find((s) => s.name === 'openai / gpt-4o') + + expect(parentSpan).toBeDefined() + expect(childSpan).toBeDefined() + + // Child span should have parent span as its parent + expect(childSpan!.parentSpanId).toBe(parentSpan!.spanContext().spanId) + + // Child span should have baggage attributes propagated + const childAttrs = childSpan!.attributes + expect(childAttrs['latitude.projectId']).toBe('456') + expect(childAttrs['latitude.commitUuid']).toBe('parent-version') + expect(childAttrs[ATTR_LATITUDE_PROMPT_PATH]).toBe('test/path') + }), + ) + + it( + 'spans created outside trace() do NOT inherit metadata', + gatewayMock.boundary(async () => { + const processorMock = new MockSpanProcessor() + + const sdk = new LatitudeTelemetry('fake-api-key', { + disableBatch: true, + processors: [processorMock], + }) + + mockRequest({ + server: gatewayMock, + method: 'post', + endpoint: '/api/v3/traces', + }) + + // First, create a traced span + await sdk.trace( + { + projectId: 789, + versionUuid: 'traced-version', + promptPath: 'traced/path', + }, + async () => { + return 'traced result' + }, + ) + + // Then, create a span OUTSIDE the trace + const completion = sdk.completion(context.active(), { + provider: 'anthropic', + model: 'claude-3', + configuration: { model: 'claude-3' }, + input: [{ role: 'user', content: 'Outside trace' }], + }) + completion.end({ + output: [{ role: 'assistant', content: 'Response' }], + tokens: { prompt: 3, cached: 0, reasoning: 0, completion: 1 }, + finishReason: 'stop', + }) + + await sdk.shutdown() + + // Get all spans + const endedSpans = processorMock.onEnd.mock.calls.map( + (call) => call[0] as ReadableSpan, + ) + + // Find the outside span + const outsideSpan = endedSpans.find( + (s) => s.name === 'anthropic / claude-3', + ) + expect(outsideSpan).toBeDefined() + + // Outside span should NOT have the trace metadata + const outsideAttrs = outsideSpan!.attributes + expect(outsideAttrs['latitude.projectId']).toBeUndefined() + expect(outsideAttrs['latitude.commitUuid']).toBeUndefined() + expect(outsideAttrs[ATTR_LATITUDE_PROMPT_PATH]).toBeUndefined() + + // Should NOT have a parent span from the trace + expect(outsideSpan!.parentSpanId).toBeUndefined() + }), + ) + + it( + 'concurrent traces do not interfere with each other', + gatewayMock.boundary(async () => { + const processorMock = new MockSpanProcessor() + + const sdk = new LatitudeTelemetry('fake-api-key', { + disableBatch: true, + processors: [processorMock], + }) + + mockRequest({ + server: gatewayMock, + method: 'post', + endpoint: '/api/v3/traces', + }) + + // Run two traces concurrently with different metadata + await Promise.all([ + sdk.trace( + { + name: 'trace-A', + projectId: 111, + promptPath: 'path/A', + }, + async () => { + // Small delay to ensure overlap + await new Promise((resolve) => setTimeout(resolve, 20)) + const completion = sdk.completion(context.active(), { + provider: 'openai', + model: 'gpt-4', + configuration: { model: 'gpt-4' }, + input: [{ role: 'user', content: 'A' }], + }) + completion.end({ + output: [{ role: 'assistant', content: 'Response A' }], + tokens: { prompt: 1, cached: 0, reasoning: 0, completion: 1 }, + finishReason: 'stop', + }) + return 'A' + }, + ), + sdk.trace( + { + name: 'trace-B', + projectId: 222, + promptPath: 'path/B', + }, + async () => { + // Small delay to ensure overlap + await new Promise((resolve) => setTimeout(resolve, 10)) + const completion = sdk.completion(context.active(), { + provider: 'anthropic', + model: 'claude', + configuration: { model: 'claude' }, + input: [{ role: 'user', content: 'B' }], + }) + completion.end({ + output: [{ role: 'assistant', content: 'Response B' }], + tokens: { prompt: 1, cached: 0, reasoning: 0, completion: 1 }, + finishReason: 'stop', + }) + return 'B' + }, + ), + ]) + + await sdk.shutdown() + + // Get all spans + const endedSpans = processorMock.onEnd.mock.calls.map( + (call) => call[0] as ReadableSpan, + ) + + // Find parent traces + const traceA = endedSpans.find((s) => s.name === 'trace-A') + const traceB = endedSpans.find((s) => s.name === 'trace-B') + + expect(traceA).toBeDefined() + expect(traceB).toBeDefined() + + // Find child completions + const childA = endedSpans.find((s) => s.name === 'openai / gpt-4') + const childB = endedSpans.find((s) => s.name === 'anthropic / claude') + + expect(childA).toBeDefined() + expect(childB).toBeDefined() + + // Child A should be parented to trace A and have A's metadata + expect(childA!.parentSpanId).toBe(traceA!.spanContext().spanId) + expect(childA!.attributes['latitude.projectId']).toBe('111') + expect(childA!.attributes[ATTR_LATITUDE_PROMPT_PATH]).toBe('path/A') + + // Child B should be parented to trace B and have B's metadata + expect(childB!.parentSpanId).toBe(traceB!.spanContext().spanId) + expect(childB!.attributes['latitude.projectId']).toBe('222') + expect(childB!.attributes[ATTR_LATITUDE_PROMPT_PATH]).toBe('path/B') + }), + ) + + it( + 'trace() propagates errors correctly', + gatewayMock.boundary(async () => { + const processorMock = new MockSpanProcessor() + + const sdk = new LatitudeTelemetry('fake-api-key', { + disableBatch: true, + processors: [processorMock], + }) + + mockRequest({ + server: gatewayMock, + method: 'post', + endpoint: '/api/v3/traces', + }) + + const testError = new Error('Test error message') + + await expect( + sdk.trace({ name: 'error-trace', projectId: 999 }, async () => { + throw testError + }), + ).rejects.toThrow('Test error message') + + await sdk.shutdown() + + // Get the trace span + const endedSpans = processorMock.onEnd.mock.calls.map( + (call) => call[0] as ReadableSpan, + ) + const errorSpan = endedSpans.find((s) => s.name === 'error-trace') + + expect(errorSpan).toBeDefined() + expect(errorSpan!.status.code).toBe(2) // SpanStatusCode.ERROR + expect(errorSpan!.status.message).toBe('Test error message') + }), + ) + + it( + 'uses HEAD_COMMIT when versionUuid is not provided', + gatewayMock.boundary(async () => { + const { bodyMock } = mockRequest({ + server: gatewayMock, + method: 'post', + endpoint: '/api/v3/traces', + }) + + const sdk = new LatitudeTelemetry('fake-api-key', { + disableBatch: true, + }) + + await sdk.trace( + { + projectId: 123, + promptPath: 'some/path', + // No versionUuid provided + }, + async () => 'result', + ) + + await sdk.shutdown() + + expect(bodyMock).toHaveBeenCalled() + const body = bodyMock.mock.calls[0]![0] + const promptSpan = body.resourceSpans[0].scopeSpans[0].spans[0] + + // Should use HEAD_COMMIT as default + expect(promptSpan.attributes).toContainEqual({ + key: 'latitude.commitUuid', + value: { stringValue: HEAD_COMMIT }, + }) + }), + ) + + it( + 'wrap() creates a reusable traced function', + gatewayMock.boundary(async () => { + const processorMock = new MockSpanProcessor() + + const sdk = new LatitudeTelemetry('fake-api-key', { + disableBatch: true, + processors: [processorMock], + }) + + mockRequest({ + server: gatewayMock, + method: 'post', + endpoint: '/api/v3/traces', + }) + + // Create a wrapped function + const tracedFn = sdk.wrap( + async (input: string) => { + const completion = sdk.completion(context.active(), { + provider: 'test', + model: 'test-model', + configuration: {}, + input: [{ role: 'user', content: input }], + }) + completion.end({ + output: [{ role: 'assistant', content: 'response' }], + tokens: { prompt: 1, cached: 0, reasoning: 0, completion: 1 }, + finishReason: 'stop', + }) + return `processed: ${input}` + }, + { + name: 'wrapped-trace', + projectId: 555, + promptPath: 'wrapped/path', + }, + ) + + // Call the wrapped function + const result = await tracedFn('test-input') + expect(result).toBe('processed: test-input') + + await sdk.shutdown() + + // Verify spans were created correctly + const endedSpans = processorMock.onEnd.mock.calls.map( + (call) => call[0] as ReadableSpan, + ) + + const parentSpan = endedSpans.find((s) => s.name === 'wrapped-trace') + const childSpan = endedSpans.find((s) => s.name === 'test / test-model') + + expect(parentSpan).toBeDefined() + expect(childSpan).toBeDefined() + expect(childSpan!.parentSpanId).toBe(parentSpan!.spanContext().spanId) + expect(childSpan!.attributes['latitude.projectId']).toBe('555') + }), + ) +}) +