diff --git a/.changeset/sunny-walls-share.md b/.changeset/sunny-walls-share.md new file mode 100644 index 000000000..692380f9e --- /dev/null +++ b/.changeset/sunny-walls-share.md @@ -0,0 +1,5 @@ +--- +"braintrust": minor +--- + +feat: Expose flag for toggling realtime mode in span fetcher diff --git a/js/src/logger.ts b/js/src/logger.ts index f315f8ed5..e9cfdcce6 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -6108,6 +6108,7 @@ export class ObjectFetcher implements AsyncIterable< // eslint-disable-next-line @typescript-eslint/no-explicit-any private mutateRecord?: (r: any) => WithTransactionId, private _internal_btql?: Record, + private _internalBrainstoreRealtime = true, ) {} public get id(): Promise { @@ -6178,7 +6179,7 @@ export class ObjectFetcher implements AsyncIterable< ...internalBtqlWithoutReservedQueryKeys, }, use_columnstore: false, - brainstore_realtime: true, + brainstore_realtime: this._internalBrainstoreRealtime, query_source: `js_sdk_object_fetcher_${this.objectType}`, ...(this.pinnedVersion !== undefined ? { diff --git a/js/src/object-fetcher.test.ts b/js/src/object-fetcher.test.ts index 931f2adcb..43d9d751b 100644 --- a/js/src/object-fetcher.test.ts +++ b/js/src/object-fetcher.test.ts @@ -31,8 +31,9 @@ class TestObjectFetcher extends ObjectFetcher { constructor( private readonly postMock: ReturnType, internalBtql?: Record, + brainstoreRealtime = true, ) { - super("dataset", undefined, undefined, internalBtql); + super("dataset", undefined, undefined, internalBtql, brainstoreRealtime); } public get id(): Promise { @@ -66,6 +67,23 @@ function getBtqlQuery( } describe("ObjectFetcher internal BTQL limit handling", () => { + test.each([ + ["default", undefined, true], + ["disabled", false, false], + ])( + "sets brainstore realtime to %s", + async (_name, inputFlag, expectedFlag) => { + const postMock = createPostMock(); + const fetcher = new TestObjectFetcher(postMock, undefined, inputFlag); + + await triggerFetch(fetcher); + + expect(postMock.mock.calls[0][1]).toMatchObject({ + brainstore_realtime: expectedFlag, + }); + }, + ); + test("preserves custom _internal_btql limit instead of default batch size", async () => { const postMock = createPostMock(); const fetcher = new TestObjectFetcher(postMock, { diff --git a/js/src/trace.test.ts b/js/src/trace.test.ts index e7a4f44f5..ee5d55fba 100644 --- a/js/src/trace.test.ts +++ b/js/src/trace.test.ts @@ -104,6 +104,31 @@ describe("CachedSpanFetcher", () => { }); }); + test("should thread brainstore realtime setting to BTQL fetches", async () => { + const post = vi.fn().mockResolvedValue({ + json: vi.fn().mockResolvedValue({ + data: [], + cursor: null, + }), + }); + const state = { + apiConn: () => ({ post }), + } as unknown as BraintrustState; + const fetcher = new CachedSpanFetcher( + "project_logs", + "project-1", + "root-1", + async () => state, + false, + ); + + await fetcher.getSpans(); + + expect(post.mock.calls[0][1]).toMatchObject({ + brainstore_realtime: false, + }); + }); + test("should fetch specific span types when filter specified", async () => { const llmSpans = [makeSpan("span-1", "llm"), makeSpan("span-2", "llm")]; diff --git a/js/src/trace.ts b/js/src/trace.ts index 5a423771e..5c0547896 100644 --- a/js/src/trace.ts +++ b/js/src/trace.ts @@ -24,6 +24,7 @@ export class SpanFetcher extends ObjectFetcher { private readonly _state: BraintrustState, private readonly spanTypeFilter?: string[], includeScorers = false, + brainstoreRealtime = true, ) { // Build the filter expression for root_span_id and optionally span_attributes.type const filterExpr = SpanFetcher.buildFilter( @@ -32,9 +33,15 @@ export class SpanFetcher extends ObjectFetcher { includeScorers, ); - super(objectType, undefined, undefined, { - filter: filterExpr, - }); + super( + objectType, + undefined, + undefined, + { + filter: filterExpr, + }, + brainstoreRealtime, + ); } private static buildFilter( @@ -138,6 +145,7 @@ export class CachedSpanFetcher { objectId: string, rootSpanId: string, getState: () => Promise, + brainstoreRealtime?: boolean, ); constructor(fetchFn: SpanFetchFn); constructor( @@ -149,6 +157,7 @@ export class CachedSpanFetcher { objectId?: string, rootSpanId?: string, getState?: () => Promise, + brainstoreRealtime = true, ) { if (typeof objectTypeOrFetchFn === "function") { // Direct fetch function injection (for testing) @@ -165,6 +174,7 @@ export class CachedSpanFetcher { state, spanType, includeScorers, + brainstoreRealtime, ); const rows: WithTransactionId[] = await fetcher.fetchedData();