Use for await, make sure readable stream works as well

js/src/langchain.ts

@@ -134,11 +134,8 @@ export class RunnableTraceable<RunInput, RunOutput> extends Runnable<
     const result = await this.invoke(input, options);
 
     if (isAsyncIterable(result)) {
-      const iterator = result[Symbol.asyncIterator]();
-      while (true) {
-        const { done, value } = await iterator.next();
-        if (done) break;
-        yield value as RunOutput;
+      for await (const item of result) {
+        yield item as RunOutput;
       }
       return;
     }

js/src/tests/traceable_langchain.test.ts

@@ -247,6 +247,41 @@ describe("to traceable", () => {
     });
   });
 
+  test("readable stream", async () => {
+    const { client, callSpy } = mockClient();
+
+    const source = RunnableTraceable.from(
+      traceable(async function (input: { text: string }) {
+        const readStream = new ReadableStream({
+          async pull(controller) {
+            for (const item of input.text.split(" ")) {
+              controller.enqueue(item);
+            }
+            controller.close();
+          },
+        });
+
+        return readStream;
+      })
+    );
+
+    const tokens: unknown[] = [];
+    for await (const chunk of await source.stream(
+      { text: "Hello world" },
+      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+      // @ts-ignore client might be of different type
+      { callbacks: [new LangChainTracer({ client })] }
+    )) {
+      tokens.push(chunk);
+    }
+
+    expect(tokens).toEqual(["Hello", "world"]);
+    expect(getAssumedTreeFromCalls(callSpy.mock.calls)).toMatchObject({
+      nodes: ["<lambda>:0"],
+      edges: [],
+    });
+  });
+
   test("async generator stream", async () => {
     const { client, callSpy } = mockClient();
     const source = RunnableTraceable.from(