diff --git a/e2e/helpers/prod-forwarding.ts b/e2e/helpers/prod-forwarding.ts index 522aedbe6..37f08d14f 100644 --- a/e2e/helpers/prod-forwarding.ts +++ b/e2e/helpers/prod-forwarding.ts @@ -31,7 +31,6 @@ export async function initializeProdForwarding(): Promise { const logger = initLogger({ apiKey, appUrl: process.env.BRAINTRUST_APP_URL, - asyncFlush: false, forceLogin: true, projectName, }); diff --git a/integrations/langchain-js/src/BraintrustCallbackHandler.ts b/integrations/langchain-js/src/BraintrustCallbackHandler.ts index d53b5ef69..56855154d 100644 --- a/integrations/langchain-js/src/BraintrustCallbackHandler.ts +++ b/integrations/langchain-js/src/BraintrustCallbackHandler.ts @@ -11,23 +11,18 @@ import { type LangChainCallbackHandlerOptions, } from "braintrust"; -type BraintrustCallbackHandlerOptions = - LangChainCallbackHandlerOptions; +type BraintrustCallbackHandlerOptions = LangChainCallbackHandlerOptions; /** * A Braintrust tracer for LangChain.js that logs LLM calls, chains, and tools. */ -export class BraintrustCallbackHandler< - IsAsyncFlush extends boolean, -> extends BaseCallbackHandler { +export class BraintrustCallbackHandler extends BaseCallbackHandler { name = "BraintrustCallbackHandler"; - private inner: BraintrustLangChainCallbackHandler; + private inner: BraintrustLangChainCallbackHandler; - constructor( - options?: Partial>, - ) { + constructor(options?: Partial) { super(); - this.inner = new BraintrustLangChainCallbackHandler(options); + this.inner = new BraintrustLangChainCallbackHandler(options); } handleLLMStart( diff --git a/integrations/openai-agents-js/src/index.ts b/integrations/openai-agents-js/src/index.ts index 394d61727..beafc3fc0 100644 --- a/integrations/openai-agents-js/src/index.ts +++ b/integrations/openai-agents-js/src/index.ts @@ -118,7 +118,7 @@ function getTimeElapsed(end?: string, start?: string): number | undefined { export class OpenAIAgentsTraceProcessor { private static readonly DEFAULT_MAX_TRACES = 10000; - private logger?: Logger; + private logger?: Logger; private maxTraces: number; private traceSpans = new Map< string, diff --git a/integrations/openai-agents-js/src/openai-agents-integration.test.ts b/integrations/openai-agents-js/src/openai-agents-integration.test.ts index 317d04b7e..71a6e5fd5 100644 --- a/integrations/openai-agents-js/src/openai-agents-integration.test.ts +++ b/integrations/openai-agents-js/src/openai-agents-integration.test.ts @@ -55,7 +55,7 @@ describe( TEST_SUITE_OPTIONS, () => { let backgroundLogger: TestBackgroundLogger; - let _logger: Logger; + let _logger: Logger; let Agent: any; let run: any; let tool: any; diff --git a/integrations/openai-agents-js/src/types.ts b/integrations/openai-agents-js/src/types.ts index 6834f5111..15aea9a6f 100644 --- a/integrations/openai-agents-js/src/types.ts +++ b/integrations/openai-agents-js/src/types.ts @@ -50,7 +50,7 @@ export type TraceMetadata = { }; export interface OpenAIAgentsTraceProcessorOptions { - logger?: Logger; + logger?: Logger; maxTraces?: number; } diff --git a/js/src/exports.ts b/js/src/exports.ts index 487ec7430..ffce08390 100644 --- a/js/src/exports.ts +++ b/js/src/exports.ts @@ -31,13 +31,11 @@ export type { InitLoggerOptions, InitOptions, LoadPromptOptions, - LogOptions, LoginOptions, Logs3OverflowInputRow, Logs3OverflowUpload, MetricSummary, ObjectMetadata, - PromiseUnless, PropagationContext, PromptRowWithId, ScoreSummary, @@ -125,7 +123,6 @@ export { withCurrent, withDataset, withExperiment, - withLogger, withParent, wrapTraced, registerOtelFlush, diff --git a/js/src/instrumentation/providers/eve-instrumentation.ts b/js/src/instrumentation/providers/eve-instrumentation.ts index 8729c6a48..610083878 100644 --- a/js/src/instrumentation/providers/eve-instrumentation.ts +++ b/js/src/instrumentation/providers/eve-instrumentation.ts @@ -585,7 +585,7 @@ class EveBridge { turn: TurnState, event: { input?: unknown; output?: unknown }, ): void { - currentLogger()?.updateSpan({ + currentLogger()?.updateSpan({ id: turn.rowId, root_span_id: turn.rootSpanId, span_id: turn.spanId, @@ -637,7 +637,7 @@ class EveBridge { if (Object.keys(totals).length === 0) { return; } - currentLogger()?.updateSpan({ + currentLogger()?.updateSpan({ id: turn.rowId, metrics: totals, root_span_id: turn.rootSpanId, diff --git a/js/src/instrumentation/providers/openai-agents-trace-processor.ts b/js/src/instrumentation/providers/openai-agents-trace-processor.ts index e3116ed58..a335f2efb 100644 --- a/js/src/instrumentation/providers/openai-agents-trace-processor.ts +++ b/js/src/instrumentation/providers/openai-agents-trace-processor.ts @@ -44,7 +44,7 @@ type TraceMetadata = { }; export interface OpenAIAgentsTraceProcessorOptions { - logger?: Logger; + logger?: Logger; maxTraces?: number; } @@ -180,7 +180,7 @@ function parseUsageMetrics(usage: unknown): Record { export class OpenAIAgentsTraceProcessor { private static readonly DEFAULT_MAX_TRACES = 10000; - private logger?: Logger; + private logger?: Logger; private maxTraces: number; private traceSpans = new Map< string, diff --git a/js/src/logger.test.ts b/js/src/logger.test.ts index 86a3a1a5c..1dffdbb49 100644 --- a/js/src/logger.test.ts +++ b/js/src/logger.test.ts @@ -1,6 +1,14 @@ /* eslint-disable @typescript-eslint/consistent-type-assertions */ -import { vi, expect, test, describe, beforeEach, afterEach } from "vitest"; +import { + vi, + expect, + expectTypeOf, + test, + describe, + beforeEach, + afterEach, +} from "vitest"; import { _exportsForTestingOnly, init, @@ -11,6 +19,7 @@ import { loadPrompt, loadParameters, wrapTraced, + traced, currentSpan, withParent, startSpan, @@ -19,6 +28,8 @@ import { deepCopyEvent, ReadonlyExperiment, renderMessageImpl, + flush, + type InitLoggerOptions, } from "./logger"; import { configureNode } from "./node/config"; @@ -31,6 +42,12 @@ import { SpanCache } from "./span-cache"; configureNode(); +type InitLoggerHasAsyncFlush = "asyncFlush" extends keyof InitLoggerOptions + ? true + : false; + +expectTypeOf().toEqualTypeOf(); + test("renderMessage with file content parts", () => { const message = { role: "user" as const, @@ -426,6 +443,57 @@ test("verify MemoryBackgroundLogger intercepts logs", async () => { _exportsForTestingOnly.clearTestBackgroundLogger(); // can go back to normal }); +test("logger and tracing APIs preserve return shapes without implicit flushing", async () => { + await _exportsForTestingOnly.simulateLoginForTests(); + const backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); + const flushSpy = vi.spyOn(backgroundLogger, "flush"); + + try { + const logger = initLogger({ + projectName: "test", + projectId: "test-project-id", + }); + + const id = logger.log({ input: "input" }); + expectTypeOf(id).toEqualTypeOf(); + + const loggerSyncResult = logger.traced(() => 42 as const); + expectTypeOf(loggerSyncResult).toEqualTypeOf<42>(); + expect(loggerSyncResult).toBe(42); + + const loggerAsyncResult = logger.traced(async () => 43); + expectTypeOf(loggerAsyncResult).toEqualTypeOf>(); + await expect(loggerAsyncResult).resolves.toBe(43); + + const tracedSyncResult = traced(() => "sync" as const); + expectTypeOf(tracedSyncResult).toEqualTypeOf<"sync">(); + expect(tracedSyncResult).toBe("sync"); + + const syncFunction = (value: number) => value + 1; + const wrappedSyncFunction = wrapTraced(syncFunction); + expectTypeOf(wrappedSyncFunction).toEqualTypeOf(); + expect(wrappedSyncFunction(1)).toBe(2); + + const asyncFunction = async (value: number) => value + 1; + const wrappedAsyncFunction = wrapTraced(asyncFunction); + expectTypeOf(wrappedAsyncFunction).toEqualTypeOf(); + await expect(wrappedAsyncFunction(2)).resolves.toBe(3); + + expect(flushSpy).not.toHaveBeenCalled(); + + await logger.flush(); + const span = logger.startSpan({ name: "explicit-flush" }); + span.end(); + await span.flush(); + await flush(); + + expect(flushSpy).toHaveBeenCalledTimes(3); + } finally { + _exportsForTestingOnly.clearTestBackgroundLogger(); + _exportsForTestingOnly.simulateLogoutForTests(); + } +}); + test("init validation", () => { expect(() => init({})).toThrow( "Must specify at least one of project or projectId", diff --git a/js/src/logger.ts b/js/src/logger.ts index fcf88d9a4..793cb96fc 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -764,9 +764,7 @@ function normalizeProxyConnUrl(proxyUrl: string): string { export class BraintrustState { public id: string; public currentExperiment: Experiment | undefined; - // Note: the value of IsAsyncFlush doesn't really matter here, since we - // (safely) dynamically cast it whenever retrieving the logger. - public currentLogger: Logger | undefined; + public currentLogger: Logger | undefined; public currentParent: IsoAsyncLocalStorage; public currentSpan: IsoAsyncLocalStorage; // Any time we re-log in, we directly update the apiConn inside the logger. @@ -1531,20 +1529,15 @@ interface OrgProjectMetadata { project: ObjectMetadata; } -export interface LinkArgs { +interface LinkArgs { org_name?: string; app_url?: string; - project_name?: string; - project_id?: string; -} - -export interface LogOptions { - asyncFlush?: IsAsyncFlush; - computeMetadataArgs?: Record; - linkArgs?: LinkArgs; } -export type PromiseUnless = B extends true ? R : Promise>; +type ProjectMetadataArgs = { + project_name?: string; + project_id?: string; +}; export interface AttachmentParams { data: string | Blob | ArrayBuffer; @@ -2546,12 +2539,11 @@ function startSpanParentArgs(args: { }; } -export class Logger implements Exportable { +export class Logger implements Exportable { private state: BraintrustState; private lazyMetadata: LazyValue; - private _asyncFlush: IsAsyncFlush | undefined; - private computeMetadataArgs: Record | undefined; - private _linkArgs: LinkArgs | undefined; + private computeMetadataArgs: ProjectMetadataArgs | undefined; + private linkArgs: LinkArgs | undefined; private lastStartTime: number; private lazyId: LazyValue; private calledStartSpan: boolean; @@ -2562,14 +2554,16 @@ export class Logger implements Exportable { constructor( state: BraintrustState, lazyMetadata: LazyValue, - logOptions: LogOptions = {}, + options: { + computeMetadataArgs?: ProjectMetadataArgs; + linkArgs?: LinkArgs; + } = {}, ) { this.lazyMetadata = lazyMetadata; - this._asyncFlush = logOptions.asyncFlush; - this.computeMetadataArgs = logOptions.computeMetadataArgs; - this._linkArgs = logOptions.linkArgs; + this.computeMetadataArgs = options.computeMetadataArgs; + this.linkArgs = options.linkArgs; this.lastStartTime = getCurrentUnixTimestamp(); - this.lazyId = new LazyValue(async () => await this.id); + this.lazyId = new LazyValue(() => this.id); this.calledStartSpan = false; this.state = state; } @@ -2599,7 +2593,7 @@ export class Logger implements Exportable { } /** - * Log a single event. The event will be batched and uploaded behind the scenes if `logOptions.asyncFlush` is true. + * Log a single event. The event will be batched and uploaded behind the scenes. Call and await {@link Logger.flush} to ensure the event has been uploaded. * * @param event The event to log. * @param event.input: (Optional) the arguments that uniquely define a user input (an arbitrary, JSON serializable object). @@ -2617,7 +2611,7 @@ export class Logger implements Exportable { public log( event: Readonly, options?: { allowConcurrentWithSpans?: boolean }, - ): PromiseUnless { + ): string { if (this.calledStartSpan && !options?.allowConcurrentWithSpans) { throw new Error( "Cannot run toplevel `log` method while using spans. To log to the span, call `logger.traced` and then log with `span.log`", @@ -2626,16 +2620,7 @@ export class Logger implements Exportable { const span = this.startSpanImpl({ startTime: this.lastStartTime, event }); this.lastStartTime = span.end(); - const ret = span.id; - type Ret = PromiseUnless; - if (this.asyncFlush === true) { - return ret as Ret; - } else { - return (async () => { - await this.flush(); - return ret; - })() as Ret; - } + return span.id; } /** @@ -2646,11 +2631,11 @@ export class Logger implements Exportable { public traced( callback: (span: Span) => R, args?: StartSpanArgs & SetCurrentArg, - ): PromiseUnless { + ): R { const { setCurrent, ...argsRest } = args ?? {}; const span = this.startSpan(argsRest); - const ret = runCatchFinally( + return runCatchFinally( () => { if (setCurrent ?? true) { return withCurrent(span, callback); @@ -2658,23 +2643,12 @@ export class Logger implements Exportable { return callback(span); } }, - (e) => { - logError(span, e); - throw e; + (error) => { + logError(span, error); + throw error; }, () => span.end(), ); - type Ret = PromiseUnless; - - if (this.asyncFlush) { - return ret as Ret; - } else { - return (async () => { - const awaitedRet = await ret; - await this.flush(); - return awaitedRet; - })() as Ret; - } } /** @@ -2773,17 +2747,13 @@ export class Logger implements Exportable { return await this.state.bgLogger().flush(); } - get asyncFlush(): IsAsyncFlush | undefined { - return this._asyncFlush; - } - /** * Return the base URL for links (e.g. https://braintrust.dev/app/my-org-name) * if we have the info, otherwise return null. * Resolution order: state -> linkArgs -> env var */ public _getLinkBaseUrl(): string | null { - return _getLinkBaseUrl(this.state, this._linkArgs); + return _getLinkBaseUrl(this.state, this.linkArgs); } /** @@ -2805,19 +2775,6 @@ export class Logger implements Exportable { } } -function castLogger( - logger: Logger | undefined, - asyncFlush?: ToB, -): Logger | undefined { - if (logger === undefined) return undefined; - if (asyncFlush !== undefined && !!asyncFlush !== !!logger.asyncFlush) { - throw new Error( - `Asserted asyncFlush setting ${asyncFlush} does not match stored logger's setting ${logger.asyncFlush}`, - ); - } - return logger as unknown as Logger; -} - export type Logs3OverflowUpload = { method: "PUT" | "POST"; signedUrl: string; @@ -4210,22 +4167,6 @@ export function withExperiment( return callback(experiment); } -/** - * @deprecated Use {@link initLogger} instead. - */ -export function withLogger( - callback: (logger: Logger) => R, - options: Readonly & SetCurrentArg> = {}, -): R { - debugLogger - .forState(options.state) - .warn( - "withLogger is deprecated and will be removed in a future version of braintrust. Simply create the logger with `initLogger`.", - ); - const logger = initLogger(options); - return callback(logger); -} - type UseOutputOption = { useOutput?: IsLegacyDataset; }; @@ -4675,13 +4616,7 @@ export function withDataset< // from arguments serialized elsewhere. async function computeLoggerMetadata( state: BraintrustState, - { - project_name, - project_id, - }: { - project_name?: string; - project_id?: string; - }, + { project_name, project_id }: ProjectMetadataArgs, ) { await state.login({}); const org_id = state.orgId!; @@ -4718,18 +4653,13 @@ async function computeLoggerMetadata( } } -type AsyncFlushArg = { - asyncFlush?: IsAsyncFlush; -}; - -export type InitLoggerOptions = FullLoginOptions & { +export type InitLoggerOptions = FullLoginOptions & { projectName?: string; projectId?: string; environment?: SpanOriginEnvironment; setCurrent?: boolean; state?: BraintrustState; - orgProjectMetadata?: OrgProjectMetadata; -} & AsyncFlushArg; +}; /** * Create a new logger in a specified project. If the project does not exist, it will be created. @@ -4737,7 +4667,6 @@ export type InitLoggerOptions = FullLoginOptions & { * @param options Additional options for configuring init(). * @param options.projectName The name of the project to log into. If unspecified, will default to the Global project. * @param options.projectId The id of the project to log into. This takes precedence over projectName if specified. - * @param options.asyncFlush If true, will log asynchronously in the background. Otherwise, will log synchronously. (true by default) * @param options.appUrl The URL of the Braintrust App. Defaults to https://www.braintrust.dev. * @param options.apiKey The API key to use. If the parameter is not specified, will try to use the `BRAINTRUST_API_KEY` environment variable. In Node.js, * if that is unset, will try the nearest `.env.braintrust` file in the current working directory or parent directories. If no API key is specified, will prompt the user to login. @@ -4747,13 +4676,10 @@ export type InitLoggerOptions = FullLoginOptions & { * @param setCurrent If true (the default), set the global current-experiment to the newly-created one. * @returns The newly created Logger. */ -export function initLogger( - options: Readonly> = {}, -) { +export function initLogger(options: Readonly = {}): Logger { const { projectName, projectId, - asyncFlush: asyncFlushArg, appUrl, apiKey, orgName, @@ -4764,9 +4690,6 @@ export function initLogger( state: stateArg, } = options || {}; - const asyncFlush = - asyncFlushArg === undefined ? (true as IsAsyncFlush) : asyncFlushArg; - const computeMetadataArgs = { project_name: projectName, project_id: projectId, @@ -4775,8 +4698,6 @@ export function initLogger( const linkArgs = { org_name: orgName, app_url: appUrl, - project_name: projectName, - project_id: projectId, }; const state = stateArg ?? _globalState; @@ -4801,13 +4722,12 @@ export function initLogger( }, ); - const ret = new Logger(state, lazyMetadata, { - asyncFlush, + const ret = new Logger(state, lazyMetadata, { computeMetadataArgs, linkArgs, }); if (options.setCurrent ?? true) { - state.currentLogger = ret as Logger; + state.currentLogger = ret; } return ret; } @@ -5484,11 +5404,9 @@ export function currentExperiment( /** * Returns the currently-active logger (set by {@link initLogger}). Returns undefined if no current logger has been set. */ -export function currentLogger( - options?: AsyncFlushArg & OptionalStateArg, -): Logger | undefined { +export function currentLogger(options?: OptionalStateArg): Logger | undefined { const state = options?.state ?? _globalState; - return castLogger(state.currentLogger, options?.asyncFlush); + return state.currentLogger; } /** @@ -5511,16 +5429,15 @@ export function currentSpan(options?: OptionalStateArg): Span { * disagreeing if state changed between calls). The state is only meaningful when * a parent slug was resolved; otherwise it is undefined. */ -function getSpanParentObjectAndPropagatedState( - options?: AsyncFlushArg & - OptionalStateArg & { parent?: string | PropagationContext }, +function getSpanParentObjectAndPropagatedState( + options?: OptionalStateArg & { parent?: string | PropagationContext }, ): { parentObject: | SpanComponentsV3 | SpanComponentsV4 | Span | Experiment - | Logger; + | Logger; propagatedState: PropagatedState | undefined; } { const state = options?.state ?? _globalState; @@ -5539,11 +5456,11 @@ function getSpanParentObjectAndPropagatedState( }; } - const experiment = currentExperiment(); + const experiment = state.currentExperiment; if (experiment) { return { parentObject: experiment, propagatedState: undefined }; } - const logger = currentLogger(options); + const logger = state.currentLogger; if (logger) { return { parentObject: logger, propagatedState: undefined }; } @@ -5557,15 +5474,9 @@ function getSpanParentObjectAndPropagatedState( * `parent` may be an exported slug string or an opaque W3C trace-context (from * {@link extractTraceContextFromHeaders}). */ -export function getSpanParentObject( - options?: AsyncFlushArg & - OptionalStateArg & { parent?: string | PropagationContext }, -): - | SpanComponentsV3 - | SpanComponentsV4 - | Span - | Experiment - | Logger { +export function getSpanParentObject( + options?: OptionalStateArg & { parent?: string | PropagationContext }, +): SpanComponentsV3 | SpanComponentsV4 | Span | Experiment | Logger { return getSpanParentObjectAndPropagatedState(options).parentObject; } @@ -6063,45 +5974,33 @@ export function logError(span: Span, error: unknown) { * * See {@link Span.traced} for full details. */ -export function traced( +export function traced( callback: (span: Span) => R, - args?: StartSpanArgs & - SetCurrentArg & - AsyncFlushArg & - OptionalStateArg, -): PromiseUnless { - const { span, isSyncFlushLogger } = startSpanAndIsLogger(args); - - const ret = runCatchFinally( + args?: StartSpanArgs & SetCurrentArg & OptionalStateArg, +): R { + const { setCurrent, ...spanArgs } = args ?? {}; + const span = startSpanImpl(spanArgs); + + return runCatchFinally( () => { - if (args?.setCurrent ?? true) { + if (setCurrent ?? true) { return withCurrent(span, callback); } else { return callback(span); } }, - (e) => { - logError(span, e); - throw e; + (error) => { + logError(span, error); + throw error; }, () => span.end(), ); - - type Ret = PromiseUnless; - - if (args?.asyncFlush === undefined || args?.asyncFlush) { - return ret as Ret; - } else { - return (async () => { - const awaitedRet = await ret; - if (isSyncFlushLogger) { - await span.flush(); - } - return awaitedRet; - })() as Ret; - } } +type WrapTracedArgs = { + noTraceIO?: boolean; +}; + /** * Check if a function is a sync generator function. * @@ -6113,7 +6012,7 @@ export function traced( * @param fn The function to check. * @returns True if the function is a sync generator function. */ -function isGeneratorFunction(fn: any): boolean { +function isGeneratorFunction(fn: unknown): boolean { return Object.prototype.toString.call(fn) === "[object GeneratorFunction]"; } @@ -6124,7 +6023,7 @@ function isGeneratorFunction(fn: any): boolean { * @param fn The function to check. * @returns True if the function is an async generator function. */ -function isAsyncGeneratorFunction(fn: any): boolean { +function isAsyncGeneratorFunction(fn: unknown): boolean { return ( Object.prototype.toString.call(fn) === "[object AsyncGeneratorFunction]" ); @@ -6135,10 +6034,13 @@ function isAsyncGeneratorFunction(fn: any): boolean { */ function wrapTracedSyncGenerator any>( fn: F, - spanArgs: any, + spanArgs: StartSpanArgs & SetCurrentArg, noTraceIO: boolean, ): F { - const wrapper = function* (this: any, ...fnArgs: Parameters) { + const wrapper = function* ( + this: ThisParameterType, + ...fnArgs: Parameters + ) { const span = startSpan(spanArgs); try { if (!noTraceIO) { @@ -6149,7 +6051,7 @@ function wrapTracedSyncGenerator any>( const maxItems = envValue !== undefined ? Number(envValue) : 1000; if (!noTraceIO && maxItems !== 0) { - let collected: any[] = []; + let collected: unknown[] = []; let truncated = false; const gen = generatorWithCurrent(span, fn.apply(this, fnArgs)); @@ -6202,10 +6104,13 @@ function wrapTracedSyncGenerator any>( */ function wrapTracedAsyncGenerator any>( fn: F, - spanArgs: any, + spanArgs: StartSpanArgs & SetCurrentArg, noTraceIO: boolean, ): F { - const wrapper = async function* (this: any, ...fnArgs: Parameters) { + const wrapper = async function* ( + this: ThisParameterType, + ...fnArgs: Parameters + ) { const span = startSpan(spanArgs); try { if (!noTraceIO) { @@ -6216,7 +6121,7 @@ function wrapTracedAsyncGenerator any>( const maxItems = envValue !== undefined ? Number(envValue) : 1000; if (!noTraceIO && maxItems !== 0) { - let collected: any[] = []; + let collected: unknown[] = []; let truncated = false; const gen = asyncGeneratorWithCurrent(span, fn.apply(this, fnArgs)); @@ -6264,10 +6169,6 @@ function wrapTracedAsyncGenerator any>( return wrapper as F; } -type WrapTracedArgs = { - noTraceIO?: boolean; -}; - /** * Wrap a function with `traced`, using the arguments as `input` and return value as `output`. * Any functions wrapped this way will automatically be traced, similar to the `@traced` decorator @@ -6293,84 +6194,52 @@ type WrapTracedArgs = { * @param args Span-level arguments (e.g. a custom name or type) to pass to `traced`. * @returns The wrapped function. */ -export function wrapTraced< - F extends (...args: any[]) => any, - IsAsyncFlush extends boolean = true, ->( +export function wrapTraced any>( fn: F, - args?: StartSpanArgs & - SetCurrentArg & - AsyncFlushArg & - WrapTracedArgs, -): IsAsyncFlush extends false - ? (...args: Parameters) => Promise>> - : F { - const spanArgs: typeof args = { + args?: StartSpanArgs & SetCurrentArg & WrapTracedArgs, +): F { + const { noTraceIO, ...argsRest } = args ?? {}; + const spanArgs: StartSpanArgs & SetCurrentArg = { name: fn.name, type: "function", - ...args, + ...argsRest, }; - const hasExplicitInput = - args && - args.event && - "input" in args.event && - args.event.input !== undefined; - const hasExplicitOutput = - args && args.event && args.event.output !== undefined; - - const noTraceIO = args?.noTraceIO || hasExplicitInput || hasExplicitOutput; + const hasExplicitInput = spanArgs.event?.input !== undefined; + const hasExplicitOutput = spanArgs.event?.output !== undefined; + + const disableGeneratorTraceIO = + !!noTraceIO || hasExplicitInput || hasExplicitOutput; // Check if the function is a generator if (isGeneratorFunction(fn)) { - return wrapTracedSyncGenerator(fn, spanArgs, !!noTraceIO); + return wrapTracedSyncGenerator(fn, spanArgs, disableGeneratorTraceIO); } if (isAsyncGeneratorFunction(fn)) { - return wrapTracedAsyncGenerator(fn, spanArgs, !!noTraceIO); + return wrapTracedAsyncGenerator(fn, spanArgs, disableGeneratorTraceIO); } - if (args?.asyncFlush) { - return ((...fnArgs: Parameters) => - traced((span) => { - if (!hasExplicitInput) { - span.log({ input: fnArgs }); - } - - const output = fn(...fnArgs); - - if (!hasExplicitOutput) { - if (output instanceof Promise) { - return (async () => { - const result = await output; - span.log({ output: result }); - return result; - })(); - } else { - span.log({ output: output }); - } - } - - return output; - }, spanArgs)) as IsAsyncFlush extends false ? never : F; - } else { - return ((...fnArgs: Parameters) => - traced(async (span) => { - if (!hasExplicitInput) { - span.log({ input: fnArgs }); - } - - const outputResult = fn(...fnArgs); + return ((...fnArgs: Parameters) => + traced((span) => { + if (!noTraceIO && !hasExplicitInput) { + span.log({ input: fnArgs }); + } - const output = await outputResult; + const output = fn(...fnArgs); - if (!hasExplicitOutput) { + if (!noTraceIO && !hasExplicitOutput) { + if (output instanceof Promise) { + return (async () => { + const result = await output; + span.log({ output: result }); + return result; + })(); + } else { span.log({ output }); } + } - return output; - }, spanArgs)) as IsAsyncFlush extends false - ? (...args: Parameters) => Promise>> - : never; - } + return output; + }, spanArgs)) as F; } /** @@ -6386,36 +6255,29 @@ export const traceable = wrapTraced; * * See {@link traced} for full details. */ -export function startSpan( - args?: StartSpanArgs & AsyncFlushArg & OptionalStateArg, -): Span { - return startSpanAndIsLogger(args).span; +export function startSpan(args?: StartSpanArgs & OptionalStateArg): Span { + return startSpanImpl(args); } /** @internal Start a span whose initial row is merged with concurrent writes. */ -export function _internalStartSpanWithInitialMerge< - IsAsyncFlush extends boolean = true, ->(args?: StartSpanArgs & AsyncFlushArg & OptionalStateArg): Span { - return startSpanAndIsLogger({ +export function _internalStartSpanWithInitialMerge( + args?: StartSpanArgs & OptionalStateArg, +): Span { + return startSpanImpl({ ...args, [INITIAL_SPAN_WRITE_AS_MERGE]: true, - } as StartSpanArgs & - AsyncFlushArg & - OptionalStateArg & - InitialSpanWriteAsMergeArg).span; + } as StartSpanArgs & OptionalStateArg & InitialSpanWriteAsMergeArg); } /** @internal Start a span with SDK-controlled context fields. */ -export function _internalStartSpanWithContext< - IsAsyncFlush extends boolean = true, ->( - args: StartSpanArgs & AsyncFlushArg & OptionalStateArg, +export function _internalStartSpanWithContext( + args: StartSpanArgs & OptionalStateArg, context: Record, ): Span { - return startSpanAndIsLogger({ + return startSpanImpl({ ...args, [INTERNAL_SPAN_CONTEXT]: context, - }).span; + }); } /** @@ -6436,23 +6298,16 @@ export function setFetch(fetch: typeof globalThis.fetch): void { _internalGetGlobalState().setFetch(fetch); } -function startSpanAndIsLogger( - args?: StartSpanArgs & - AsyncFlushArg & - OptionalStateArg & - InternalSpanContextArg, -): { span: Span; isSyncFlushLogger: boolean } { +function startSpanImpl( + args?: StartSpanArgs & OptionalStateArg & InternalSpanContextArg, +): Span { const state = args?.state ?? _globalState; // Resolve the parent object and any forwarded W3C state in one pass, so we // don't re-normalize `parent` (which could disagree if the active // logger/experiment changed between calls). const { parentObject, propagatedState } = - getSpanParentObjectAndPropagatedState({ - asyncFlush: args?.asyncFlush, - parent: args?.parent, - state, - }); + getSpanParentObjectAndPropagatedState(args); if ( parentObject instanceof SpanComponentsV3 || @@ -6472,7 +6327,7 @@ function startSpanAndIsLogger( // The parent object/state are already resolved from `parent` above; drop // the raw `parent` so it isn't re-normalized. const { parent: _ignoredParent, ...spanArgs } = args ?? {}; - const span = new SpanImpl({ + return new SpanImpl({ state, ...spanArgs, parentObjectType: parentObject.data.object_type, @@ -6490,21 +6345,8 @@ function startSpanAndIsLogger( | undefined), propagatedState, }); - return { - span, - isSyncFlushLogger: - parentObject.data.object_type === SpanObjectTypeV3.PROJECT_LOGS && - // Since there's no parent logger here, we're free to choose the async flush - // behavior, and therefore propagate along whatever we get from the arguments - args?.asyncFlush === false, - }; } else { - const span = parentObject.startSpan(args); - return { - span, - isSyncFlushLogger: - parentObject.kind === "logger" && parentObject.asyncFlush === false, - }; + return parentObject.startSpan(args); } } @@ -7197,7 +7039,7 @@ export class Experiment this.lazyMetadata = lazyMetadata; this.dataset = dataset; this.lastStartTime = getCurrentUnixTimestamp(); - this.lazyId = new LazyValue(async () => await this.id); + this.lazyId = new LazyValue(() => this.id); this.calledStartSpan = false; this.state = state; } @@ -7298,7 +7140,7 @@ export class Experiment const { setCurrent, ...argsRest } = args ?? {}; const span = this.startSpan(argsRest); - const ret = runCatchFinally( + return runCatchFinally( () => { if (setCurrent ?? true) { return withCurrent(span, callback); @@ -7306,14 +7148,12 @@ export class Experiment return callback(span); } }, - (e) => { - logError(span, e); - throw e; + (error) => { + logError(span, error); + throw error; }, () => span.end(), ); - - return ret as R; } /** @@ -7941,9 +7781,9 @@ export class SpanImpl implements Span { return callback(span); } }, - (e) => { - logError(span, e); - throw e; + (error) => { + logError(span, error); + throw error; }, () => span.end(), ); diff --git a/js/src/vendor-sdk-types/langchain.ts b/js/src/vendor-sdk-types/langchain.ts index a4ffdc909..fd795b8b5 100644 --- a/js/src/vendor-sdk-types/langchain.ts +++ b/js/src/vendor-sdk-types/langchain.ts @@ -27,10 +27,10 @@ export type LangChainCallbackManagerConfigureArgs = [ options?: unknown, ]; -export type LangChainCallbackHandlerOptions = { +export type LangChainCallbackHandlerOptions = { debug: boolean; excludeMetadataProps: RegExp; - logger?: Logger | Span; + logger?: Logger | Span; parent?: Span | (() => Span); }; diff --git a/js/src/wrappers/langchain/callback-handler.ts b/js/src/wrappers/langchain/callback-handler.ts index a0143b338..51b073173 100644 --- a/js/src/wrappers/langchain/callback-handler.ts +++ b/js/src/wrappers/langchain/callback-handler.ts @@ -15,22 +15,18 @@ import type { export const BRAINTRUST_LANGCHAIN_CALLBACK_HANDLER_NAME = "BraintrustCallbackHandler"; -export class BraintrustLangChainCallbackHandler< - IsAsyncFlush extends boolean = true, -> { +export class BraintrustLangChainCallbackHandler { name = BRAINTRUST_LANGCHAIN_CALLBACK_HANDLER_NAME; private spans = new Map(); private skippedRuns = new Set(); private parent?: Span | (() => Span); private rootRunId?: string; - private options: LangChainCallbackHandlerOptions; + private options: LangChainCallbackHandlerOptions; private startTimes = new Map(); private firstTokenTimes = new Map(); private ttftMs = new Map(); - constructor( - options?: Partial>, - ) { + constructor(options?: Partial) { this.parent = options?.parent; this.options = { debug: options?.debug ?? false, diff --git a/js/src/wrappers/test_oai_attachments.test.ts b/js/src/wrappers/test_oai_attachments.test.ts index 6fcc0ec72..443086900 100644 --- a/js/src/wrappers/test_oai_attachments.test.ts +++ b/js/src/wrappers/test_oai_attachments.test.ts @@ -11,7 +11,6 @@ describe("OpenAI attachment processing", () => { initLogger({ projectName: PROJECT_NAME, apiKey: "test-key", - asyncFlush: false, }); }); diff --git a/js/tests/api-compatibility/api-compatibility.test.ts b/js/tests/api-compatibility/api-compatibility.test.ts index 2fd785357..3941302be 100644 --- a/js/tests/api-compatibility/api-compatibility.test.ts +++ b/js/tests/api-compatibility/api-compatibility.test.ts @@ -2223,8 +2223,8 @@ describe("isUnionTypeWidening", () => { describe("areInterfaceSignaturesCompatible", () => { test("should allow adding optional fields to interface", () => { - const oldInterface = `export interface LogOptions { asyncFlush?: IsAsyncFlush; computeMetadataArgs?: Record; }`; - const newInterface = `export interface LogOptions { asyncFlush?: IsAsyncFlush; computeMetadataArgs?: Record; linkArgs?: LinkArgs; }`; + const oldInterface = `export interface FeatureOptions { enabled?: Enabled; metadata?: Record; }`; + const newInterface = `export interface FeatureOptions { enabled?: Enabled; metadata?: Record; link?: string; }`; const result = areInterfaceSignaturesCompatible(oldInterface, newInterface); expect(result).toBe(true); @@ -2239,40 +2239,40 @@ describe("areInterfaceSignaturesCompatible", () => { }); test("should reject removing fields from interface", () => { - const oldInterface = `export interface LogOptions { asyncFlush?: IsAsyncFlush; computeMetadataArgs?: Record; }`; - const newInterface = `export interface LogOptions { asyncFlush?: IsAsyncFlush; }`; + const oldInterface = `export interface FeatureOptions { enabled?: Enabled; metadata?: Record; }`; + const newInterface = `export interface FeatureOptions { enabled?: Enabled; }`; const result = areInterfaceSignaturesCompatible(oldInterface, newInterface); expect(result).toBe(false); }); test("should reject adding required fields to interface", () => { - const oldInterface = `export interface LogOptions { asyncFlush?: IsAsyncFlush; }`; - const newInterface = `export interface LogOptions { asyncFlush?: IsAsyncFlush; requiredField: string; }`; + const oldInterface = `export interface FeatureOptions { enabled?: Enabled; }`; + const newInterface = `export interface FeatureOptions { enabled?: Enabled; requiredField: string; }`; const result = areInterfaceSignaturesCompatible(oldInterface, newInterface); expect(result).toBe(false); }); test("should reject changing field types in interface", () => { - const oldInterface = `export interface LogOptions { asyncFlush?: IsAsyncFlush; computeMetadataArgs?: Record; }`; - const newInterface = `export interface LogOptions { asyncFlush?: IsAsyncFlush; computeMetadataArgs?: string; }`; + const oldInterface = `export interface FeatureOptions { enabled?: Enabled; metadata?: Record; }`; + const newInterface = `export interface FeatureOptions { enabled?: Enabled; metadata?: string; }`; const result = areInterfaceSignaturesCompatible(oldInterface, newInterface); expect(result).toBe(false); }); test("should allow making required field optional", () => { - const oldInterface = `export interface LogOptions { asyncFlush: IsAsyncFlush; }`; - const newInterface = `export interface LogOptions { asyncFlush?: IsAsyncFlush; }`; + const oldInterface = `export interface FeatureOptions { enabled: Enabled; }`; + const newInterface = `export interface FeatureOptions { enabled?: Enabled; }`; const result = areInterfaceSignaturesCompatible(oldInterface, newInterface); expect(result).toBe(true); }); test("should reject making optional field required", () => { - const oldInterface = `export interface LogOptions { asyncFlush?: IsAsyncFlush; }`; - const newInterface = `export interface LogOptions { asyncFlush: IsAsyncFlush; }`; + const oldInterface = `export interface FeatureOptions { enabled?: Enabled; }`; + const newInterface = `export interface FeatureOptions { enabled: Enabled; }`; const result = areInterfaceSignaturesCompatible(oldInterface, newInterface); expect(result).toBe(false); @@ -2577,8 +2577,8 @@ describe("areFunctionSignaturesCompatible", () => { test("should allow widening a field inside an options-object parameter", () => { // Real-world case: getSpanParentObject's options param widens its nested // `parent?` field from `string` to `string | PropagationContext`. - const oldFn = `export function getSpanParentObject(options?: AsyncFlushArg & OptionalStateArg & { parent?: string; }): Span`; - const newFn = `export function getSpanParentObject(options?: AsyncFlushArg & OptionalStateArg & { parent?: string | PropagationContext; }): Span`; + const oldFn = `export function getSpanParentObject(options?: ModeArg & OptionalStateArg & { parent?: string; }): Span`; + const newFn = `export function getSpanParentObject(options?: ModeArg & OptionalStateArg & { parent?: string | PropagationContext; }): Span`; expect(areFunctionSignaturesCompatible(oldFn, newFn)).toBe(true); });