fix(core): Await default client batches when awaiting all callbacks, set default traced project from LangSmith

langchain-core/src/singletons/callbacks.ts

@@ -2,6 +2,7 @@
 
 import PQueueMod from "p-queue";
 import { getGlobalAsyncLocalStorageInstance } from "./async_local_storage/globals.js";
+import { getDefaultLangChainClientSingleton } from "./tracer.js";
 
 let queue: typeof import("p-queue")["default"]["prototype"];
 
@@ -60,6 +61,10 @@ export async function consumeCallback<T>(
  * Waits for all promises in the queue to resolve. If the queue is
  * undefined, it immediately resolves a promise.
  */
-export function awaitAllCallbacks(): Promise<void> {
-  return typeof queue !== "undefined" ? queue.onIdle() : Promise.resolve();
+export async function awaitAllCallbacks(): Promise<void> {
+  const defaultClient = getDefaultLangChainClientSingleton();
+  await Promise.all([
+    typeof queue !== "undefined" ? queue.onIdle() : Promise.resolve(),
+    defaultClient.awaitPendingTraceBatches(),
+  ]);
 }

langchain-core/src/tracers/tests/tracer.test.ts

@@ -1,15 +1,23 @@
 import { test, expect, jest } from "@jest/globals";
 import * as uuid from "uuid";
+import { Client } from "langsmith";
 import { Serialized } from "../../load/serializable.js";
 import { Document } from "../../documents/document.js";
 import { Run } from "../base.js";
 import { HumanMessage } from "../../messages/index.js";
 import { FakeTracer } from "../../utils/testing/index.js";
+import { setDefaultLangChainClientSingleton } from "../../singletons/tracer.js";
+import { RunnableLambda } from "../../runnables/index.js";
+import { awaitAllCallbacks } from "../../singletons/callbacks.js";
 
 const _DATE = 1620000000000;
 
 Date.now = jest.fn(() => _DATE);
 
+afterEach(() => {
+  setDefaultLangChainClientSingleton(new Client());
+});
+
 const serialized: Serialized = {
   lc: 1,
   type: "constructor",
@@ -410,3 +418,29 @@
   await tracer.handleLLMEnd({ generations: [[]] }, llmRunId3);
   expect(tracer.runs.length).toBe(2);
 });
+
+test("Test tracer payload snapshots for run create and update", async () => {
+  const client = new Client();
+  (client as any).multipartIngestRuns = jest.fn(async () => {
+    await new Promise((resolve) => setTimeout(resolve, 500));
+    return Promise.resolve();
+  });
+  setDefaultLangChainClientSingleton(client);
+  const parentRunnable = RunnableLambda.from(async (input: string) => {
+    const childRunnable = RunnableLambda.from(async (childInput: string) => {
+      return `processed: ${childInput}`;
+    });
+
+    const result = await childRunnable.invoke(input);
+    return `parent: ${result}`;
+  });
+  await parentRunnable.invoke("test input");
+
+  const beforeAwaitTime = new Date();
+  await awaitAllCallbacks();
+  const afterAwaitTime = new Date();
+  expect(afterAwaitTime.getTime() - beforeAwaitTime.getTime()).toBeGreaterThan(
+    500
+  );
+  expect((client as any).multipartIngestRuns).toHaveBeenCalled();
+});