langsmith 0.1.23 → 0.1.24

package/dist/evaluation/evaluate_comparative.js

@@ -0,0 +1,195 @@
+import { v4 as uuid4, validate } from "uuid";
+import { Client } from "../index.js";
+import { shuffle } from "../utils/shuffle.js";
+import { AsyncCaller } from "../utils/async_caller.js";
+import pRetry from "p-retry";
+import { getCurrentRunTree, traceable } from "../traceable.js";
+function isExperimentResultsList(value) {
+    return value.some((x) => typeof x !== "string");
+}
+async function loadExperiment(client, experiment) {
+    const value = typeof experiment === "string" ? experiment : experiment.experimentName;
+    return client.readProject(validate(value) ? { projectId: value } : { projectName: value });
+}
+async function loadTraces(client, experiment, options) {
+    const executionOrder = options.loadNested ? undefined : 1;
+    const runs = await client.listRuns(validate(experiment)
+        ? { projectId: experiment, executionOrder }
+        : { projectName: experiment, executionOrder });
+    const treeMap = {};
+    const runIdMap = {};
+    const results = [];
+    for await (const run of runs) {
+        if (run.parent_run_id != null) {
+            treeMap[run.parent_run_id] ??= [];
+            treeMap[run.parent_run_id].push(run);
+        }
+        else {
+            results.push(run);
+        }
+        runIdMap[run.id] = run;
+    }
+    for (const [parentRunId, childRuns] of Object.entries(treeMap)) {
+        const parentRun = runIdMap[parentRunId];
+        parentRun.child_runs = childRuns.sort((a, b) => {
+            if (a.dotted_order == null || b.dotted_order == null)
+                return 0;
+            return a.dotted_order.localeCompare(b.dotted_order);
+        });
+    }
+    return results;
+}
+export async function evaluateComparative(experiments, options) {
+    if (experiments.length < 2) {
+        throw new Error("Comparative evaluation requires at least 2 experiments.");
+    }
+    if (!options.evaluators.length) {
+        throw new Error("At least one evaluator is required for comparative evaluation.");
+    }
+    if (options.maxConcurrency && options.maxConcurrency < 0) {
+        throw new Error("maxConcurrency must be a positive number.");
+    }
+    const client = options.client ?? new Client();
+    const resolvedExperiments = await Promise.all(experiments);
+    const projects = await (() => {
+        if (!isExperimentResultsList(resolvedExperiments)) {
+            return Promise.all(resolvedExperiments.map((experiment) => loadExperiment(client, experiment)));
+        }
+        // if we know the number of runs beforehand, check if the
+        // number of runs in the project matches the expected number of runs
+        return Promise.all(resolvedExperiments.map((experiment) => pRetry(async () => {
+            const project = await loadExperiment(client, experiment);
+            if (project.run_count !== experiment?.results.length) {
+                throw new Error("Experiment is missing runs. Retrying.");
+            }
+            return project;
+        }, { factor: 2, minTimeout: 1000, retries: 10 })));
+    })();
+    if (new Set(projects.map((p) => p.reference_dataset_id)).size > 1) {
+        throw new Error("All experiments must have the same reference dataset.");
+    }
+    const referenceDatasetId = projects.at(0)?.reference_dataset_id;
+    if (!referenceDatasetId) {
+        throw new Error("Reference dataset is required for comparative evaluation.");
+    }
+    if (new Set(projects.map((p) => p.extra?.metadata?.dataset_version)).size > 1) {
+        console.warn("Detected multiple dataset versions used by experiments, which may lead to inaccurate results.");
+    }
+    const datasetVersion = projects.at(0)?.extra?.metadata?.dataset_version;
+    const id = uuid4();
+    const experimentName = (() => {
+        if (!options.experimentPrefix) {
+            const names = projects
+                .map((p) => p.name)
+                .filter(Boolean)
+                .join(" vs. ");
+            return `${names}-${uuid4().slice(0, 4)}`;
+        }
+        return `${options.experimentPrefix}-${uuid4().slice(0, 4)}`;
+    })();
+    // TODO: add URL to the comparative experiment
+    console.log(`Starting pairwise evaluation of: ${experimentName}`);
+    const comparativeExperiment = await client.createComparativeExperiment({
+        id,
+        name: experimentName,
+        experimentIds: projects.map((p) => p.id),
+        description: options.description,
+        metadata: options.metadata,
+        referenceDatasetId: projects.at(0)?.reference_dataset_id,
+    });
+    const viewUrl = await (async () => {
+        const projectId = projects.at(0)?.id ?? projects.at(1)?.id;
+        const datasetId = comparativeExperiment?.reference_dataset_id;
+        if (projectId && datasetId) {
+            const hostUrl = (await client.getProjectUrl({ projectId }))
+                .split("/projects/p/")
+                .at(0);
+            const result = new URL(`${hostUrl}/datasets/${datasetId}/compare`);
+            result.searchParams.set("selectedSessions", projects.map((p) => p.id).join(","));
+            result.searchParams.set("comparativeExperiment", comparativeExperiment.id);
+            return result.toString();
+        }
+        return null;
+    })();
+    if (viewUrl != null) {
+        console.log(`View results at: ${viewUrl}`);
+    }
+    const experimentRuns = await Promise.all(projects.map((p) => loadTraces(client, p.id, { loadNested: !!options.loadNested })));
+    let exampleIdsIntersect;
+    for (const runs of experimentRuns) {
+        const exampleIdsSet = new Set(runs
+            .map((r) => r.reference_example_id)
+            .filter((x) => x != null));
+        if (!exampleIdsIntersect) {
+            exampleIdsIntersect = exampleIdsSet;
+        }
+        else {
+            exampleIdsIntersect = new Set([...exampleIdsIntersect].filter((x) => exampleIdsSet.has(x)));
+        }
+    }
+    const exampleIds = [...(exampleIdsIntersect ?? [])];
+    if (!exampleIds.length) {
+        throw new Error("No examples found in common between experiments.");
+    }
+    const exampleMap = {};
+    for (let start = 0; start < exampleIds.length; start += 99) {
+        const exampleIdsChunk = exampleIds.slice(start, start + 99);
+        for await (const example of client.listExamples({
+            datasetId: referenceDatasetId,
+            exampleIds: exampleIdsChunk,
+            asOf: datasetVersion,
+        })) {
+            exampleMap[example.id] = example;
+        }
+    }
+    const runMapByExampleId = {};
+    for (const runs of experimentRuns) {
+        for (const run of runs) {
+            if (run.reference_example_id == null ||
+                !exampleIds.includes(run.reference_example_id)) {
+                continue;
+            }
+            runMapByExampleId[run.reference_example_id] ??= [];
+            runMapByExampleId[run.reference_example_id].push(run);
+        }
+    }
+    const caller = new AsyncCaller({ maxConcurrency: options.maxConcurrency });
+    async function evaluateAndSubmitFeedback(runs, example, evaluator) {
+        const expectedRunIds = new Set(runs.map((r) => r.id));
+        const result = await evaluator(options.randomizeOrder ? shuffle(runs) : runs, example);
+        for (const [runId, score] of Object.entries(result.scores)) {
+            // validate if the run id
+            if (!expectedRunIds.has(runId)) {
+                throw new Error(`Returning an invalid run id ${runId} from evaluator.`);
+            }
+            await client.createFeedback(runId, result.key, {
+                score,
+                sourceRunId: result.source_run_id,
+                comparativeExperimentId: comparativeExperiment.id,
+            });
+        }
+        return result;
+    }
+    const tracedEvaluators = options.evaluators.map((evaluator) => traceable(async (runs, example) => {
+        const evaluatorRun = getCurrentRunTree();
+        const result = await evaluator(runs, example);
+        // sanitise the payload before sending to LangSmith
+        evaluatorRun.inputs = { runs: runs, example: example };
+        evaluatorRun.outputs = result;
+        return {
+            ...result,
+            source_run_id: result.source_run_id ?? evaluatorRun.id,
+        };
+    }, {
+        project_name: "evaluators",
+        name: evaluator.name || "evaluator",
+    }));
+    const promises = Object.entries(runMapByExampleId).flatMap(([exampleId, runs]) => {
+        const example = exampleMap[exampleId];
+        if (!example)
+            throw new Error(`Example ${exampleId} not found.`);
+        return tracedEvaluators.map((evaluator) => caller.call(evaluateAndSubmitFeedback, runs, exampleMap[exampleId], evaluator));
+    });
+    const results = await Promise.all(promises);
+    return { experimentName, results };
+}

package/dist/evaluation/evaluator.cjs

@@ -14,11 +14,10 @@ class DynamicRunEvaluator {
             writable: true,
             value: void 0
         });
-        const wrappedFunc = (input) => {
-            const runAndExample = input.langSmithRunAndExample;
-            return evaluator(...Object.values(runAndExample));
-        };
-        this.func = wrappedFunc;
+        this.func = ((input) => {
+            const { run, example } = input.langSmithRunAndExample;
+            return evaluator(run, example);
+        });
     }
     coerceEvaluationResults(results, sourceRunId) {
         if ("results" in results) {
@@ -57,17 +56,15 @@ class DynamicRunEvaluator {
         if ("session_id" in run) {
             metadata["experiment"] = run.session_id;
         }
-        const wrappedTraceableFunc = (0, traceable_js_1.wrapFunctionAndEnsureTraceable)(this.func, options || {}, "evaluator");
+        if (typeof this.func !== "function") {
+            throw new Error("Target must be runnable function");
+        }
+        const wrappedTraceableFunc = (0, traceable_js_1.traceable)(this.func, { project_name: "evaluators", name: "evaluator", ...options });
+        const result = (await wrappedTraceableFunc(
         // Pass data via `langSmithRunAndExample` key to avoid conflicts with other
         // inputs. This key is extracted in the wrapped function, with `run` and
         // `example` passed to evaluator function as arguments.
-        const langSmithRunAndExample = {
-            run,
-            example,
-        };
-        const result = (await wrappedTraceableFunc({ langSmithRunAndExample }, {
-            metadata,
-        }));
+        { langSmithRunAndExample: { run, example } }, { metadata }));
         // Check the one required property of EvaluationResult since 'instanceof' is not possible
         if ("key" in result) {
             if (!result.sourceRunId) {

package/dist/evaluation/evaluator.js

@@ -1,5 +1,5 @@
 import { v4 as uuidv4 } from "uuid";
-import { wrapFunctionAndEnsureTraceable } from "../traceable.js";
+import { traceable } from "../traceable.js";
 /**
  * Wraps an evaluator function + implements the RunEvaluator interface.
  */
@@ -11,11 +11,10 @@ export class DynamicRunEvaluator {
             writable: true,
             value: void 0
         });
-        const wrappedFunc = (input) => {
-            const runAndExample = input.langSmithRunAndExample;
-            return evaluator(...Object.values(runAndExample));
-        };
-        this.func = wrappedFunc;
+        this.func = ((input) => {
+            const { run, example } = input.langSmithRunAndExample;
+            return evaluator(run, example);
+        });
     }
     coerceEvaluationResults(results, sourceRunId) {
         if ("results" in results) {
@@ -54,17 +53,15 @@ export class DynamicRunEvaluator {
         if ("session_id" in run) {
             metadata["experiment"] = run.session_id;
         }
-        const wrappedTraceableFunc = wrapFunctionAndEnsureTraceable(this.func, options || {}, "evaluator");
+        if (typeof this.func !== "function") {
+            throw new Error("Target must be runnable function");
+        }
+        const wrappedTraceableFunc = traceable(this.func, { project_name: "evaluators", name: "evaluator", ...options });
+        const result = (await wrappedTraceableFunc(
         // Pass data via `langSmithRunAndExample` key to avoid conflicts with other
         // inputs. This key is extracted in the wrapped function, with `run` and
         // `example` passed to evaluator function as arguments.
-        const langSmithRunAndExample = {
-            run,
-            example,
-        };
-        const result = (await wrappedTraceableFunc({ langSmithRunAndExample }, {
-            metadata,
-        }));
+        { langSmithRunAndExample: { run, example } }, { metadata }));
         // Check the one required property of EvaluationResult since 'instanceof' is not possible
         if ("key" in result) {
             if (!result.sourceRunId) {

package/dist/evaluation/index.cjs

@@ -1,7 +1,9 @@
 "use strict";
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.evaluate = exports.StringEvaluator = void 0;
+exports.evaluateComparative = exports.evaluate = exports.StringEvaluator = void 0;
 var string_evaluator_js_1 = require("./string_evaluator.cjs");
 Object.defineProperty(exports, "StringEvaluator", { enumerable: true, get: function () { return string_evaluator_js_1.StringEvaluator; } });
 var _runner_js_1 = require("./_runner.cjs");
 Object.defineProperty(exports, "evaluate", { enumerable: true, get: function () { return _runner_js_1.evaluate; } });
+var evaluate_comparative_js_1 = require("./evaluate_comparative.cjs");
+Object.defineProperty(exports, "evaluateComparative", { enumerable: true, get: function () { return evaluate_comparative_js_1.evaluateComparative; } });

package/dist/evaluation/index.d.ts

@@ -1,3 +1,4 @@
 export { RunEvaluator, EvaluationResult } from "./evaluator.js";
 export { StringEvaluator, GradingFunctionParams, GradingFunctionResult, } from "./string_evaluator.js";
 export { evaluate, type EvaluateOptions } from "./_runner.js";
+export { evaluateComparative } from "./evaluate_comparative.js";

package/dist/evaluation/index.js

@@ -1,2 +1,3 @@
 export { StringEvaluator, } from "./string_evaluator.js";
 export { evaluate } from "./_runner.js";
+export { evaluateComparative } from "./evaluate_comparative.js";

package/dist/index.cjs

@@ -6,4 +6,4 @@ Object.defineProperty(exports, "Client", { enumerable: true, get: function () {
 var run_trees_js_1 = require("./run_trees.cjs");
 Object.defineProperty(exports, "RunTree", { enumerable: true, get: function () { return run_trees_js_1.RunTree; } });
 // Update using yarn bump-version
-exports.__version__ = "0.1.23";
+exports.__version__ = "0.1.24";

package/dist/index.d.ts

@@ -1,4 +1,4 @@
 export { Client } from "./client.js";
 export type { Dataset, Example, TracerSession, Run, Feedback, } from "./schemas.js";
 export { RunTree, type RunTreeConfig } from "./run_trees.js";
-export declare const __version__ = "0.1.23";
+export declare const __version__ = "0.1.24";

package/dist/index.js

@@ -1,4 +1,4 @@
 export { Client } from "./client.js";
 export { RunTree } from "./run_trees.js";
 // Update using yarn bump-version
-export const __version__ = "0.1.23";
+export const __version__ = "0.1.24";

package/dist/run_trees.cjs

@@ -339,6 +339,7 @@ class RunTree {
         const runUpdate = {
             end_time: this.end_time,
             error: this.error,
+            inputs: this.inputs,
             outputs: this.outputs,
             parent_run_id: this.parent_run?.id,
             reference_example_id: this.reference_example_id,

package/dist/run_trees.js

@@ -312,6 +312,7 @@ export class RunTree {
         const runUpdate = {
             end_time: this.end_time,
             error: this.error,
+            inputs: this.inputs,
             outputs: this.outputs,
             parent_run_id: this.parent_run?.id,
             reference_example_id: this.reference_example_id,

package/dist/schemas.d.ts

@@ -276,4 +276,21 @@ export interface DatasetDiffInfo {
     examples_added: string[];
     examples_removed: string[];
 }
+export interface ComparisonEvaluationResult {
+    key: string;
+    scores: Record<string, ScoreType>;
+    source_run_id?: string;
+}
+export interface ComparativeExperiment {
+    id: string;
+    name: string;
+    description: string;
+    tenant_id: string;
+    created_at: string;
+    modified_at: string;
+    reference_dataset_id: string;
+    extra?: Record<string, unknown>;
+    experiments_info?: Array<Record<string, unknown>>;
+    feedback_stats?: Record<string, unknown>;
+}
 export {};

package/dist/traceable.cjs

@@ -1,7 +1,7 @@
 "use strict";
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.wrapFunctionAndEnsureTraceable = exports.isTraceableFunction = exports.getCurrentRunTree = exports.traceable = exports.ROOT = void 0;
-const async_hooks_1 = require("async_hooks");
+const node_async_hooks_1 = require("node:async_hooks");
 const run_trees_js_1 = require("./run_trees.cjs");
 const env_js_1 = require("./utils/env.cjs");
 function isPromiseMethod(x) {
@@ -10,12 +10,28 @@ function isPromiseMethod(x) {
     }
     return false;
 }
-const asyncLocalStorage = new async_hooks_1.AsyncLocalStorage();
+const asyncLocalStorage = new node_async_hooks_1.AsyncLocalStorage();
 exports.ROOT = Symbol("langsmith:traceable:root");
 const isAsyncIterable = (x) => x != null &&
     typeof x === "object" &&
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     typeof x[Symbol.asyncIterator] === "function";
+const GeneratorFunction = function* () { }.constructor;
+const isIteratorLike = (x) => x != null &&
+    typeof x === "object" &&
+    "next" in x &&
+    typeof x.next === "function";
+const isGenerator = (x) => 
+// eslint-disable-next-line no-instanceof/no-instanceof
+x != null && typeof x === "function" && x instanceof GeneratorFunction;
+const isThenable = (x) => x != null &&
+    typeof x === "object" &&
+    "then" in x &&
+    typeof x.then === "function";
+const isReadableStream = (x) => x != null &&
+    typeof x === "object" &&
+    "getReader" in x &&
+    typeof x.getReader === "function";
 const tracingIsEnabled = (tracingEnabled) => {
     if (tracingEnabled !== undefined) {
         return tracingEnabled;
@@ -55,6 +71,141 @@ const getTracingRunTree = (runTree, inputs) => {
     runTree.inputs = handleRunInputs(inputs);
     return runTree;
 };
+// idea: store the state of the promise outside
+// but only when the promise is "consumed"
+const getSerializablePromise = (arg) => {
+    const proxyState = { current: undefined };
+    const promiseProxy = new Proxy(arg, {
+        get(target, prop, receiver) {
+            if (prop === "then") {
+                const boundThen = arg[prop].bind(arg);
+                return (resolve, reject = (x) => {
+                    throw x;
+                }) => {
+                    return boundThen((value) => {
+                        proxyState.current = ["resolve", value];
+                        return resolve(value);
+                    }, (error) => {
+                        proxyState.current = ["reject", error];
+                        return reject(error);
+                    });
+                };
+            }
+            if (prop === "catch") {
+                const boundCatch = arg[prop].bind(arg);
+                return (reject) => {
+                    return boundCatch((error) => {
+                        proxyState.current = ["reject", error];
+                        return reject(error);
+                    });
+                };
+            }
+            if (prop === "toJSON") {
+                return () => {
+                    if (!proxyState.current)
+                        return undefined;
+                    const [type, value] = proxyState.current ?? [];
+                    if (type === "resolve")
+                        return value;
+                    return { error: value };
+                };
+            }
+            return Reflect.get(target, prop, receiver);
+        },
+    });
+    return promiseProxy;
+};
+const convertSerializableArg = (arg) => {
+    if (isReadableStream(arg)) {
+        const proxyState = [];
+        const transform = new TransformStream({
+            start: () => void 0,
+            transform: (chunk, controller) => {
+                proxyState.push(chunk);
+                controller.enqueue(chunk);
+            },
+            flush: () => void 0,
+        });
+        const pipeThrough = arg.pipeThrough(transform);
+        Object.assign(pipeThrough, { toJSON: () => proxyState });
+        return pipeThrough;
+    }
+    if (isAsyncIterable(arg)) {
+        const proxyState = { current: [] };
+        return new Proxy(arg, {
+            get(target, prop, receiver) {
+                if (prop === Symbol.asyncIterator) {
+                    return () => {
+                        const boundIterator = arg[Symbol.asyncIterator].bind(arg);
+                        const iterator = boundIterator();
+                        return new Proxy(iterator, {
+                            get(target, prop, receiver) {
+                                if (prop === "next" || prop === "return" || prop === "throw") {
+                                    const bound = iterator.next.bind(iterator);
+                                    return (...args) => {
+                                        // @ts-expect-error TS cannot infer the argument types for the bound function
+                                        const wrapped = getSerializablePromise(bound(...args));
+                                        proxyState.current.push(wrapped);
+                                        return wrapped;
+                                    };
+                                }
+                                if (prop === "return" || prop === "throw") {
+                                    return iterator.next.bind(iterator);
+                                }
+                                return Reflect.get(target, prop, receiver);
+                            },
+                        });
+                    };
+                }
+                if (prop === "toJSON") {
+                    return () => {
+                        const onlyNexts = proxyState.current;
+                        const serialized = onlyNexts.map((next) => next.toJSON());
+                        const chunks = serialized.reduce((memo, next) => {
+                            if (next?.value)
+                                memo.push(next.value);
+                            return memo;
+                        }, []);
+                        return chunks;
+                    };
+                }
+                return Reflect.get(target, prop, receiver);
+            },
+        });
+    }
+    if (!Array.isArray(arg) && isIteratorLike(arg)) {
+        const proxyState = [];
+        return new Proxy(arg, {
+            get(target, prop, receiver) {
+                if (prop === "next" || prop === "return" || prop === "throw") {
+                    const bound = arg[prop]?.bind(arg);
+                    return (...args) => {
+                        // @ts-expect-error TS cannot infer the argument types for the bound function
+                        const next = bound?.(...args);
+                        if (next != null)
+                            proxyState.push(next);
+                        return next;
+                    };
+                }
+                if (prop === "toJSON") {
+                    return () => {
+                        const chunks = proxyState.reduce((memo, next) => {
+                            if (next.value)
+                                memo.push(next.value);
+                            return memo;
+                        }, []);
+                        return chunks;
+                    };
+                }
+                return Reflect.get(target, prop, receiver);
+            },
+        });
+    }
+    if (isThenable(arg)) {
+        return getSerializablePromise(arg);
+    }
+    return arg;
+};
 /**
  * Higher-order function that takes function as input and returns a
  * "TraceableFunction" - a wrapped version of the input that
@@ -118,8 +269,13 @@ function traceable(wrappedFunc, config) {
                 ...runTreeConfig,
             };
         }
+        // TODO: deal with possible nested promises and async iterables
+        const processedArgs = args;
+        for (let i = 0; i < processedArgs.length; i++) {
+            processedArgs[i] = convertSerializableArg(processedArgs[i]);
+        }
         const [currentRunTree, rawInputs] = (() => {
-            const [firstArg, ...restArgs] = args;
+            const [firstArg, ...restArgs] = processedArgs;
             // used for handoff between LangChain.JS and traceable functions
             if ((0, run_trees_js_1.isRunnableConfigLike)(firstArg)) {
                 return [
@@ -147,12 +303,12 @@ function traceable(wrappedFunc, config) {
             const prevRunFromStore = asyncLocalStorage.getStore();
             if (prevRunFromStore) {
                 return [
-                    getTracingRunTree(prevRunFromStore.createChild(ensuredConfig), args),
-                    args,
+                    getTracingRunTree(prevRunFromStore.createChild(ensuredConfig), processedArgs),
+                    processedArgs,
                 ];
             }
-            const currentRunTree = getTracingRunTree(new run_trees_js_1.RunTree(ensuredConfig), args);
-            return [currentRunTree, args];
+            const currentRunTree = getTracingRunTree(new run_trees_js_1.RunTree(ensuredConfig), processedArgs);
+            return [currentRunTree, processedArgs];
         })();
         return asyncLocalStorage.run(currentRunTree, () => {
             const postRunPromise = currentRunTree?.postRun();
@@ -208,6 +364,17 @@ function traceable(wrappedFunc, config) {
                 await postRunPromise;
                 await currentRunTree?.patchRun();
             }
+            function gatherAll(iterator) {
+                const chunks = [];
+                // eslint-disable-next-line no-constant-condition
+                while (true) {
+                    const next = iterator.next();
+                    chunks.push(next);
+                    if (next.done)
+                        break;
+                }
+                return chunks;
+            }
             let returnValue;
             try {
                 returnValue = wrappedFunc(...rawInputs);
@@ -216,25 +383,40 @@ function traceable(wrappedFunc, config) {
                 returnValue = Promise.reject(err);
             }
             if (isAsyncIterable(returnValue)) {
-                const snapshot = async_hooks_1.AsyncLocalStorage.snapshot();
+                const snapshot = node_async_hooks_1.AsyncLocalStorage.snapshot();
                 return wrapAsyncGeneratorForTracing(returnValue, snapshot);
             }
             const tracedPromise = new Promise((resolve, reject) => {
                 Promise.resolve(returnValue)
                     .then(async (rawOutput) => {
                     if (isAsyncIterable(rawOutput)) {
-                        const snapshot = async_hooks_1.AsyncLocalStorage.snapshot();
+                        const snapshot = node_async_hooks_1.AsyncLocalStorage.snapshot();
                         return resolve(wrapAsyncGeneratorForTracing(rawOutput, snapshot));
                     }
-                    else {
-                        try {
-                            await currentRunTree?.end(handleRunOutputs(rawOutput));
-                            await handleEnd();
-                        }
-                        finally {
-                            // eslint-disable-next-line no-unsafe-finally
-                            return rawOutput;
-                        }
+                    if (isGenerator(wrappedFunc) && isIteratorLike(rawOutput)) {
+                        const chunks = gatherAll(rawOutput);
+                        await currentRunTree?.end(handleRunOutputs(await handleChunks(chunks.reduce((memo, { value, done }) => {
+                            if (!done || typeof value !== "undefined") {
+                                memo.push(value);
+                            }
+                            return memo;
+                        }, []))));
+                        await handleEnd();
+                        return (function* () {
+                            for (const ret of chunks) {
+                                if (ret.done)
+                                    return ret.value;
+                                yield ret.value;
+                            }
+                        })();
+                    }
+                    try {
+                        await currentRunTree?.end(handleRunOutputs(rawOutput));
+                        await handleEnd();
+                    }
+                    finally {
+                        // eslint-disable-next-line no-unsafe-finally
+                        return rawOutput;
                     }
                 }, async (error) => {
                     await currentRunTree?.end(undefined, String(error));
@@ -287,11 +469,15 @@ function isTraceableFunction(x
 }
 exports.isTraceableFunction = isTraceableFunction;
 function isKVMap(x) {
-    return (typeof x === "object" &&
-        x != null &&
-        !Array.isArray(x) &&
-        // eslint-disable-next-line no-instanceof/no-instanceof
-        !(x instanceof Date));
+    if (typeof x !== "object" || x == null) {
+        return false;
+    }
+    const prototype = Object.getPrototypeOf(x);
+    return ((prototype === null ||
+        prototype === Object.prototype ||
+        Object.getPrototypeOf(prototype) === null) &&
+        !(Symbol.toStringTag in x) &&
+        !(Symbol.iterator in x));
 }
 function wrapFunctionAndEnsureTraceable(target, options, name = "target") {
     if (typeof target === "function") {