Fix ordering of recorded runs, buffer runs

langchain-ai · Oct 22, 2024 · 997204c · 997204c
1 parent aae789a
commit 997204c
Show file tree

Hide file tree

Showing 3 changed files with 421 additions and 317 deletions.
diff --git a/js/src/tests/utils/iterator.ts b/js/src/tests/utils/iterator.ts
@@ -0,0 +1,7 @@
+export async function gatherIterator<T>(
+  i: AsyncIterable<T> | Promise<AsyncIterable<T>>
+): Promise<Array<T>> {
+  const out: T[] = [];
+  for await (const item of await i) out.push(item);
+  return out;
+}
diff --git a/js/src/tests/vercel_exporter.int.test.ts b/js/src/tests/vercel_exporter.int.test.ts
@@ -12,6 +12,7 @@ import {
   embedMany,
 } from "ai";
 import { tool } from "ai";
+import { gatherIterator } from "./utils/iterator.js";
 
 const telemetrySettings = {
   isEnabled: true,
@@ -22,19 +23,11 @@ const telemetrySettings = {
   },
 };
 
-test.concurrent("generateText", async () => {
-  const traceExporter = new LangSmithAISDKExporter();
-  const sdk = new NodeSDK({ traceExporter });
-  sdk.start();
-
-  function getOrders(userId: string) {
-    return `User ${userId} has the following orders: 1`;
-  }
-
-  function getTrackingInformation(orderId: string) {
-    return `Here is the tracking information for ${orderId}`;
-  }
+const traceExporter = new LangSmithAISDKExporter();
+const sdk = new NodeSDK({ traceExporter });
+sdk.start();
 
+test("generateText", async () => {
   await generateText({
     model: openai("gpt-4o-mini"),
     messages: [
@@ -47,34 +40,48 @@ test.concurrent("generateText", async () => {
       listOrders: tool({
         description: "list all orders",
         parameters: z.object({ userId: z.string() }),
-        execute: async ({ userId }) => getOrders(userId),
+        execute: async ({ userId }) =>
+          `User ${userId} has the following orders: 1`,
       }),
       viewTrackingInformation: tool({
         description: "view tracking information for a specific order",
         parameters: z.object({ orderId: z.string() }),
-        execute: async ({ orderId }) => getTrackingInformation(orderId),
+        execute: async ({ orderId }) =>
+          `Here is the tracking information for ${orderId}`,
       }),
     },
     experimental_telemetry: telemetrySettings,
     maxSteps: 10,
   });
 
-  await sdk.shutdown();
+  await traceExporter.forceFlush?.();
 });
 
-test.concurrent("streamText", async () => {
-  const traceExporter = new LangSmithAISDKExporter();
-  const sdk = new NodeSDK({ traceExporter });
-  sdk.start();
-
-  function getOrders(userId: string) {
-    return `User ${userId} has the following orders: 1`;
-  }
+test("generateText with image", async () => {
+  await generateText({
+    model: openai("gpt-4o-mini"),
+    messages: [
+      {
+        role: "user",
+        content: [
+          {
+            type: "text",
+            text: "What's in this picture?",
+          },
+          {
+            type: "image",
+            image: new URL("https://picsum.photos/200/300"),
+          },
+        ],
+      },
+    ],
+    experimental_telemetry: telemetrySettings,
+  });
 
-  function getTrackingInformation(orderId: string) {
-    return `Here is the tracking information for ${orderId}`;
-  }
+  await traceExporter.forceFlush?.();
+});
 
+test("streamText", async () => {
   const result = await streamText({
     model: openai("gpt-4o-mini"),
     messages: [
@@ -87,30 +94,25 @@ test.concurrent("streamText", async () => {
       listOrders: tool({
         description: "list all orders",
         parameters: z.object({ userId: z.string() }),
-        execute: async ({ userId }) => getOrders(userId),
+        execute: async ({ userId }) =>
+          `User ${userId} has the following orders: 1`,
       }),
       viewTrackingInformation: tool({
         description: "view tracking information for a specific order",
         parameters: z.object({ orderId: z.string() }),
-        execute: async ({ orderId }) => getTrackingInformation(orderId),
+        execute: async ({ orderId }) =>
+          `Here is the tracking information for ${orderId}`,
       }),
     },
     experimental_telemetry: { isEnabled: true },
     maxSteps: 10,
   });
 
-  for await (const _stream of result.fullStream) {
-    // consume
-  }
-
-  await sdk.shutdown();
+  await gatherIterator(result.fullStream);
+  await traceExporter.forceFlush?.();
 });
 
-test.concurrent("generateObject", async () => {
-  const traceExporter = new LangSmithAISDKExporter();
-  const sdk = new NodeSDK({ traceExporter });
-  sdk.start();
-
+test("generateObject", async () => {
   await generateObject({
     model: openai("gpt-4o-mini", { structuredOutputs: true }),
     schema: z.object({
@@ -123,14 +125,10 @@ test.concurrent("generateObject", async () => {
     experimental_telemetry: telemetrySettings,
   });
 
-  await sdk.shutdown();
+  await traceExporter.forceFlush?.();
 });
 
-test.concurrent("streamObject", async () => {
-  const traceExporter = new LangSmithAISDKExporter();
-  const sdk = new NodeSDK({ traceExporter });
-  sdk.start();
-
+test("streamObject", async () => {
   const result = await streamObject({
     model: openai("gpt-4o-mini", { structuredOutputs: true }),
     schema: z.object({
@@ -143,32 +141,21 @@ test.concurrent("streamObject", async () => {
     experimental_telemetry: telemetrySettings,
   });
 
-  for await (const _partialObject of result.partialObjectStream) {
-    // pass
-  }
-
-  await sdk.shutdown();
+  await gatherIterator(result.partialObjectStream);
+  await traceExporter.forceFlush?.();
 });
 
-test.concurrent("embed", async () => {
-  const traceExporter = new LangSmithAISDKExporter();
-  const sdk = new NodeSDK({ traceExporter });
-  sdk.start();
-
+test("embed", async () => {
   await embed({
     model: openai.embedding("text-embedding-3-small"),
     value: "prague castle at sunset",
     experimental_telemetry: telemetrySettings,
   });
 
-  await sdk.shutdown();
+  await traceExporter.forceFlush?.();
 });
 
-test.concurrent("embedMany", async () => {
-  const traceExporter = new LangSmithAISDKExporter();
-  const sdk = new NodeSDK({ traceExporter });
-  sdk.start();
-
+test("embedMany", async () => {
   await embedMany({
     model: openai.embedding("text-embedding-3-small"),
     values: [
@@ -179,33 +166,9 @@ test.concurrent("embedMany", async () => {
     experimental_telemetry: telemetrySettings,
   });
 
-  await sdk.shutdown();
+  await traceExporter.forceFlush?.();
 });
 
-test.concurrent("generateText with image", async () => {
-  const traceExporter = new LangSmithAISDKExporter();
-  const sdk = new NodeSDK({ traceExporter });
-  sdk.start();
-
-  await generateText({
-    model: openai("gpt-4o-mini"),
-    messages: [
-      {
-        role: "user",
-        content: [
-          {
-            type: "text",
-            text: "What's in this picture?",
-          },
-          {
-            type: "image",
-            image: new URL("https://picsum.photos/200/300"),
-          },
-        ],
-      },
-    ],
-    experimental_telemetry: telemetrySettings,
-  });
-
+afterAll(async () => {
   await sdk.shutdown();
 });