langsmith 0.1.20 → 0.1.21

package/dist/client.cjs

@@ -148,7 +148,7 @@ class Queue {
 }
 exports.Queue = Queue;
 // 20 MB
-exports.DEFAULT_BATCH_SIZE_LIMIT_BYTES = 20971520;
+exports.DEFAULT_BATCH_SIZE_LIMIT_BYTES = 20_971_520;
 class Client {
     constructor(config = {}) {
         Object.defineProperty(this, "apiKey", {
@@ -276,7 +276,7 @@ class Client {
         this.apiUrl = trimQuotes(config.apiUrl ?? defaultConfig.apiUrl) ?? "";
         this.apiKey = trimQuotes(config.apiKey ?? defaultConfig.apiKey);
         this.webUrl = trimQuotes(config.webUrl ?? defaultConfig.webUrl);
-        this.timeout_ms = config.timeout_ms ?? 12000;
+        this.timeout_ms = config.timeout_ms ?? 12_000;
         this.caller = new async_caller_js_1.AsyncCaller(config.callerOptions ?? {});
         this.batchIngestCaller = new async_caller_js_1.AsyncCaller({
             ...(config.callerOptions ?? {}),

package/dist/client.js

@@ -121,7 +121,7 @@ export class Queue {
     }
 }
 // 20 MB
-export const DEFAULT_BATCH_SIZE_LIMIT_BYTES = 20971520;
+export const DEFAULT_BATCH_SIZE_LIMIT_BYTES = 20_971_520;
 export class Client {
     constructor(config = {}) {
         Object.defineProperty(this, "apiKey", {
@@ -249,7 +249,7 @@ export class Client {
         this.apiUrl = trimQuotes(config.apiUrl ?? defaultConfig.apiUrl) ?? "";
         this.apiKey = trimQuotes(config.apiKey ?? defaultConfig.apiKey);
         this.webUrl = trimQuotes(config.webUrl ?? defaultConfig.webUrl);
-        this.timeout_ms = config.timeout_ms ?? 12000;
+        this.timeout_ms = config.timeout_ms ?? 12_000;
         this.caller = new AsyncCaller(config.callerOptions ?? {});
         this.batchIngestCaller = new AsyncCaller({
             ...(config.callerOptions ?? {}),

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.20";
+exports.__version__ = "0.1.21";

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.20";
+export declare const __version__ = "0.1.21";

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.20";
+export const __version__ = "0.1.21";

package/dist/run_trees.cjs

@@ -232,7 +232,7 @@ class RunTree {
             extra: {},
         };
     }
-    async createChild(config) {
+    createChild(config) {
         const child = new RunTree({
             ...config,
             parent_run: this,
@@ -243,9 +243,9 @@ class RunTree {
         return child;
     }
     async end(outputs, error, endTime = Date.now()) {
-        this.outputs = outputs;
-        this.error = error;
-        this.end_time = endTime;
+        this.outputs = this.outputs ?? outputs;
+        this.error = this.error ?? error;
+        this.end_time = this.end_time ?? endTime;
     }
     async _convertToCreate(run, excludeChildRuns = true) {
         const runExtra = run.extra ?? {};

package/dist/run_trees.d.ts

@@ -65,7 +65,7 @@ export declare class RunTree implements BaseRun {
         metadata?: KVMap;
     }): RunTree;
     private static getDefaultConfig;
-    createChild(config: RunTreeConfig): Promise<RunTree>;
+    createChild(config: RunTreeConfig): RunTree;
     end(outputs?: KVMap, error?: string, endTime?: number): Promise<void>;
     private _convertToCreate;
     postRun(excludeChildRuns?: boolean): Promise<void>;

package/dist/run_trees.js

@@ -205,7 +205,7 @@ export class RunTree {
             extra: {},
         };
     }
-    async createChild(config) {
+    createChild(config) {
         const child = new RunTree({
             ...config,
             parent_run: this,
@@ -216,9 +216,9 @@ export class RunTree {
         return child;
     }
     async end(outputs, error, endTime = Date.now()) {
-        this.outputs = outputs;
-        this.error = error;
-        this.end_time = endTime;
+        this.outputs = this.outputs ?? outputs;
+        this.error = this.error ?? error;
+        this.end_time = this.end_time ?? endTime;
     }
     async _convertToCreate(run, excludeChildRuns = true) {
         const runExtra = run.extra ?? {};

package/dist/traceable.cjs

@@ -4,6 +4,12 @@ exports.isTraceableFunction = exports.getCurrentRunTree = exports.traceable = vo
 const async_hooks_1 = require("async_hooks");
 const run_trees_js_1 = require("./run_trees.cjs");
 const env_js_1 = require("./utils/env.cjs");
+function isPromiseMethod(x) {
+    if (x === "then" || x === "catch" || x === "finally") {
+        return true;
+    }
+    return false;
+}
 const asyncLocalStorage = new async_hooks_1.AsyncLocalStorage();
 const isAsyncIterable = (x) => x != null &&
     typeof x === "object" &&
@@ -33,14 +39,55 @@ const getTracingRunTree = (runTree) => {
  */
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 function traceable(wrappedFunc, config) {
-    const { aggregator, ...runTreeConfig } = config ?? {};
-    const traceableFunc = async (...args) => {
+    const { aggregator, argsConfigPath, ...runTreeConfig } = config ?? {};
+    const traceableFunc = (...args) => {
         let currentRunTree;
         let rawInputs;
-        const ensuredConfig = {
-            name: wrappedFunc.name || "<lambda>",
-            ...runTreeConfig,
-        };
+        let ensuredConfig;
+        try {
+            let runtimeConfig;
+            if (argsConfigPath) {
+                const [index, path] = argsConfigPath;
+                if (index === args.length - 1 && !path) {
+                    runtimeConfig = args.pop();
+                }
+                else if (index <= args.length &&
+                    typeof args[index] === "object" &&
+                    args[index] !== null) {
+                    if (path) {
+                        const { [path]: extracted, ...rest } = args[index];
+                        runtimeConfig = extracted;
+                        args[index] = rest;
+                    }
+                    else {
+                        runtimeConfig = args[index];
+                        args.splice(index, 1);
+                    }
+                }
+            }
+            ensuredConfig = {
+                name: wrappedFunc.name || "<lambda>",
+                ...runTreeConfig,
+                ...runtimeConfig,
+                tags: [
+                    ...new Set([
+                        ...(runTreeConfig?.tags ?? []),
+                        ...(runtimeConfig?.tags ?? []),
+                    ]),
+                ],
+                metadata: {
+                    ...runTreeConfig?.metadata,
+                    ...runtimeConfig?.metadata,
+                },
+            };
+        }
+        catch (err) {
+            console.warn(`Failed to extract runtime config from args for ${runTreeConfig?.name ?? wrappedFunc.name}`, err);
+            ensuredConfig = {
+                name: wrappedFunc.name || "<lambda>",
+                ...runTreeConfig,
+            };
+        }
         const previousRunTree = asyncLocalStorage.getStore();
         if ((0, run_trees_js_1.isRunTree)(args[0])) {
             currentRunTree = args[0];
@@ -51,7 +98,7 @@ function traceable(wrappedFunc, config) {
             rawInputs = args.slice(1);
         }
         else if (previousRunTree !== undefined) {
-            currentRunTree = await previousRunTree.createChild(ensuredConfig);
+            currentRunTree = previousRunTree.createChild(ensuredConfig);
             rawInputs = args;
         }
         else {
@@ -76,74 +123,145 @@ function traceable(wrappedFunc, config) {
         if (currentRunTree) {
             currentRunTree.inputs = inputs;
         }
-        const initialOutputs = currentRunTree?.outputs;
-        const initialError = currentRunTree?.error;
-        await currentRunTree?.postRun();
-        return new Promise((resolve, reject) => {
-            void asyncLocalStorage.run(currentRunTree, async () => {
-                try {
-                    const rawOutput = await wrappedFunc(...rawInputs);
+        return asyncLocalStorage.run(currentRunTree, () => {
+            const postRunPromise = currentRunTree?.postRun();
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            let returnValue;
+            try {
+                returnValue = wrappedFunc(...rawInputs);
+            }
+            catch (err) {
+                returnValue = Promise.reject(err);
+            }
+            if (isAsyncIterable(returnValue)) {
+                // eslint-disable-next-line no-inner-declarations
+                async function* wrapOutputForTracing() {
+                    let finished = false;
+                    const chunks = [];
+                    try {
+                        for await (const chunk of returnValue) {
+                            chunks.push(chunk);
+                            yield chunk;
+                        }
+                        finished = true;
+                    }
+                    catch (e) {
+                        await currentRunTree?.end(undefined, String(e));
+                        throw e;
+                    }
+                    finally {
+                        if (!finished) {
+                            await currentRunTree?.end(undefined, "Cancelled");
+                        }
+                        let finalOutputs;
+                        if (aggregator !== undefined) {
+                            try {
+                                finalOutputs = await aggregator(chunks);
+                            }
+                            catch (e) {
+                                console.error(`[ERROR]: LangSmith aggregation failed: `, e);
+                                finalOutputs = chunks;
+                            }
+                        }
+                        else {
+                            finalOutputs = chunks;
+                        }
+                        if (typeof finalOutputs === "object" &&
+                            !Array.isArray(finalOutputs)) {
+                            await currentRunTree?.end(finalOutputs);
+                        }
+                        else {
+                            await currentRunTree?.end({ outputs: finalOutputs });
+                        }
+                        await postRunPromise;
+                        await currentRunTree?.patchRun();
+                    }
+                }
+                return wrapOutputForTracing();
+            }
+            const tracedPromise = new Promise((resolve, reject) => {
+                Promise.resolve(returnValue)
+                    .then(
+                // eslint-disable-next-line @typescript-eslint/no-explicit-any
+                async (rawOutput) => {
                     if (isAsyncIterable(rawOutput)) {
                         // eslint-disable-next-line no-inner-declarations
                         async function* wrapOutputForTracing() {
+                            let finished = false;
                             const chunks = [];
-                            // TypeScript thinks this is unsafe
-                            for await (const chunk of rawOutput) {
-                                chunks.push(chunk);
-                                yield chunk;
+                            try {
+                                // TypeScript thinks this is unsafe
+                                for await (const chunk of rawOutput) {
+                                    chunks.push(chunk);
+                                    yield chunk;
+                                }
+                                finished = true;
+                            }
+                            catch (e) {
+                                await currentRunTree?.end(undefined, String(e));
+                                throw e;
                             }
-                            let finalOutputs;
-                            if (aggregator !== undefined) {
-                                try {
-                                    finalOutputs = await aggregator(chunks);
+                            finally {
+                                if (!finished) {
+                                    await currentRunTree?.end(undefined, "Cancelled");
                                 }
-                                catch (e) {
-                                    console.error(`[ERROR]: LangSmith aggregation failed: `, e);
+                                let finalOutputs;
+                                if (aggregator !== undefined) {
+                                    try {
+                                        finalOutputs = await aggregator(chunks);
+                                    }
+                                    catch (e) {
+                                        console.error(`[ERROR]: LangSmith aggregation failed: `, e);
+                                        finalOutputs = chunks;
+                                    }
+                                }
+                                else {
                                     finalOutputs = chunks;
                                 }
+                                if (typeof finalOutputs === "object" &&
+                                    !Array.isArray(finalOutputs)) {
+                                    await currentRunTree?.end(finalOutputs);
+                                }
+                                else {
+                                    await currentRunTree?.end({ outputs: finalOutputs });
+                                }
+                                await postRunPromise;
+                                await currentRunTree?.patchRun();
                             }
-                            else {
-                                finalOutputs = chunks;
-                            }
-                            if (typeof finalOutputs === "object" &&
-                                !Array.isArray(finalOutputs)) {
-                                await currentRunTree?.end(finalOutputs);
-                            }
-                            else {
-                                await currentRunTree?.end({ outputs: finalOutputs });
-                            }
-                            await currentRunTree?.patchRun();
                         }
                         return resolve(wrapOutputForTracing());
                     }
                     else {
-                        const outputs = isKVMap(rawOutput)
-                            ? rawOutput
-                            : { outputs: rawOutput };
-                        if (initialOutputs === currentRunTree?.outputs) {
-                            await currentRunTree?.end(outputs);
-                        }
-                        else {
-                            if (currentRunTree !== undefined) {
-                                currentRunTree.end_time = Date.now();
-                            }
+                        try {
+                            await currentRunTree?.end(isKVMap(rawOutput) ? rawOutput : { outputs: rawOutput });
+                            await postRunPromise;
+                            await currentRunTree?.patchRun();
                         }
-                        await currentRunTree?.patchRun();
-                        return resolve(rawOutput);
-                    }
-                }
-                catch (error) {
-                    if (initialError === currentRunTree?.error) {
-                        await currentRunTree?.end(initialOutputs, String(error));
-                    }
-                    else {
-                        if (currentRunTree !== undefined) {
-                            currentRunTree.end_time = Date.now();
+                        finally {
+                            // eslint-disable-next-line no-unsafe-finally
+                            return rawOutput;
                         }
                     }
+                }, 
+                // eslint-disable-next-line @typescript-eslint/no-explicit-any
+                async (error) => {
+                    await currentRunTree?.end(undefined, String(error));
+                    await postRunPromise;
                     await currentRunTree?.patchRun();
-                    reject(error);
-                }
+                    throw error;
+                })
+                    .then(resolve, reject);
+            });
+            if (typeof returnValue !== "object" || returnValue === null) {
+                return tracedPromise;
+            }
+            return new Proxy(returnValue, {
+                get(target, prop, receiver) {
+                    if (isPromiseMethod(prop)) {
+                        return tracedPromise[prop].bind(tracedPromise);
+                    }
+                    return Reflect.get(target, prop, receiver);
+                },
             });
         });
     };

package/dist/traceable.d.ts

@@ -1,12 +1,13 @@
 import { RunTree, RunTreeConfig, RunnableConfigLike } from "./run_trees.js";
 export type RunTreeLike = RunTree;
+type SmartPromise<T> = T extends AsyncGenerator ? T : T extends Promise<unknown> ? T : Promise<T>;
 type WrapArgReturnPair<Pair> = Pair extends [
     infer Args extends any[],
     infer Return
 ] ? {
-    (...args: Args): Promise<Return>;
-    (...args: [runTree: RunTreeLike, ...rest: Args]): Promise<Return>;
-    (...args: [config: RunnableConfigLike, ...rest: Args]): Promise<Return>;
+    (...args: Args): SmartPromise<Return>;
+    (...args: [runTree: RunTreeLike, ...rest: Args]): SmartPromise<Return>;
+    (...args: [config: RunnableConfigLike, ...rest: Args]): SmartPromise<Return>;
 } : never;
 type UnionToIntersection<U> = (U extends any ? (x: U) => void : never) extends (x: infer I) => void ? I : never;
 export type TraceableFunction<Func extends (...args: any[]) => any> = Func extends {
@@ -46,6 +47,7 @@ export type TraceableFunction<Func extends (...args: any[]) => any> = Func exten
  */
 export declare function traceable<Func extends (...args: any[]) => any>(wrappedFunc: Func, config?: Partial<RunTreeConfig> & {
     aggregator?: (args: any[]) => any;
+    argsConfigPath?: [number] | [number, string];
 }): TraceableFunction<Func>;
 /**
  * Return the current run tree from within a traceable-wrapped function.

package/dist/traceable.js

@@ -1,6 +1,12 @@
 import { AsyncLocalStorage } from "async_hooks";
 import { RunTree, isRunTree, isRunnableConfigLike, } from "./run_trees.js";
 import { getEnvironmentVariable } from "./utils/env.js";
+function isPromiseMethod(x) {
+    if (x === "then" || x === "catch" || x === "finally") {
+        return true;
+    }
+    return false;
+}
 const asyncLocalStorage = new AsyncLocalStorage();
 const isAsyncIterable = (x) => x != null &&
     typeof x === "object" &&
@@ -30,14 +36,55 @@ const getTracingRunTree = (runTree) => {
  */
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 export function traceable(wrappedFunc, config) {
-    const { aggregator, ...runTreeConfig } = config ?? {};
-    const traceableFunc = async (...args) => {
+    const { aggregator, argsConfigPath, ...runTreeConfig } = config ?? {};
+    const traceableFunc = (...args) => {
         let currentRunTree;
         let rawInputs;
-        const ensuredConfig = {
-            name: wrappedFunc.name || "<lambda>",
-            ...runTreeConfig,
-        };
+        let ensuredConfig;
+        try {
+            let runtimeConfig;
+            if (argsConfigPath) {
+                const [index, path] = argsConfigPath;
+                if (index === args.length - 1 && !path) {
+                    runtimeConfig = args.pop();
+                }
+                else if (index <= args.length &&
+                    typeof args[index] === "object" &&
+                    args[index] !== null) {
+                    if (path) {
+                        const { [path]: extracted, ...rest } = args[index];
+                        runtimeConfig = extracted;
+                        args[index] = rest;
+                    }
+                    else {
+                        runtimeConfig = args[index];
+                        args.splice(index, 1);
+                    }
+                }
+            }
+            ensuredConfig = {
+                name: wrappedFunc.name || "<lambda>",
+                ...runTreeConfig,
+                ...runtimeConfig,
+                tags: [
+                    ...new Set([
+                        ...(runTreeConfig?.tags ?? []),
+                        ...(runtimeConfig?.tags ?? []),
+                    ]),
+                ],
+                metadata: {
+                    ...runTreeConfig?.metadata,
+                    ...runtimeConfig?.metadata,
+                },
+            };
+        }
+        catch (err) {
+            console.warn(`Failed to extract runtime config from args for ${runTreeConfig?.name ?? wrappedFunc.name}`, err);
+            ensuredConfig = {
+                name: wrappedFunc.name || "<lambda>",
+                ...runTreeConfig,
+            };
+        }
         const previousRunTree = asyncLocalStorage.getStore();
         if (isRunTree(args[0])) {
             currentRunTree = args[0];
@@ -48,7 +95,7 @@ export function traceable(wrappedFunc, config) {
             rawInputs = args.slice(1);
         }
         else if (previousRunTree !== undefined) {
-            currentRunTree = await previousRunTree.createChild(ensuredConfig);
+            currentRunTree = previousRunTree.createChild(ensuredConfig);
             rawInputs = args;
         }
         else {
@@ -73,74 +120,145 @@ export function traceable(wrappedFunc, config) {
         if (currentRunTree) {
             currentRunTree.inputs = inputs;
         }
-        const initialOutputs = currentRunTree?.outputs;
-        const initialError = currentRunTree?.error;
-        await currentRunTree?.postRun();
-        return new Promise((resolve, reject) => {
-            void asyncLocalStorage.run(currentRunTree, async () => {
-                try {
-                    const rawOutput = await wrappedFunc(...rawInputs);
+        return asyncLocalStorage.run(currentRunTree, () => {
+            const postRunPromise = currentRunTree?.postRun();
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            let returnValue;
+            try {
+                returnValue = wrappedFunc(...rawInputs);
+            }
+            catch (err) {
+                returnValue = Promise.reject(err);
+            }
+            if (isAsyncIterable(returnValue)) {
+                // eslint-disable-next-line no-inner-declarations
+                async function* wrapOutputForTracing() {
+                    let finished = false;
+                    const chunks = [];
+                    try {
+                        for await (const chunk of returnValue) {
+                            chunks.push(chunk);
+                            yield chunk;
+                        }
+                        finished = true;
+                    }
+                    catch (e) {
+                        await currentRunTree?.end(undefined, String(e));
+                        throw e;
+                    }
+                    finally {
+                        if (!finished) {
+                            await currentRunTree?.end(undefined, "Cancelled");
+                        }
+                        let finalOutputs;
+                        if (aggregator !== undefined) {
+                            try {
+                                finalOutputs = await aggregator(chunks);
+                            }
+                            catch (e) {
+                                console.error(`[ERROR]: LangSmith aggregation failed: `, e);
+                                finalOutputs = chunks;
+                            }
+                        }
+                        else {
+                            finalOutputs = chunks;
+                        }
+                        if (typeof finalOutputs === "object" &&
+                            !Array.isArray(finalOutputs)) {
+                            await currentRunTree?.end(finalOutputs);
+                        }
+                        else {
+                            await currentRunTree?.end({ outputs: finalOutputs });
+                        }
+                        await postRunPromise;
+                        await currentRunTree?.patchRun();
+                    }
+                }
+                return wrapOutputForTracing();
+            }
+            const tracedPromise = new Promise((resolve, reject) => {
+                Promise.resolve(returnValue)
+                    .then(
+                // eslint-disable-next-line @typescript-eslint/no-explicit-any
+                async (rawOutput) => {
                     if (isAsyncIterable(rawOutput)) {
                         // eslint-disable-next-line no-inner-declarations
                         async function* wrapOutputForTracing() {
+                            let finished = false;
                             const chunks = [];
-                            // TypeScript thinks this is unsafe
-                            for await (const chunk of rawOutput) {
-                                chunks.push(chunk);
-                                yield chunk;
+                            try {
+                                // TypeScript thinks this is unsafe
+                                for await (const chunk of rawOutput) {
+                                    chunks.push(chunk);
+                                    yield chunk;
+                                }
+                                finished = true;
+                            }
+                            catch (e) {
+                                await currentRunTree?.end(undefined, String(e));
+                                throw e;
                             }
-                            let finalOutputs;
-                            if (aggregator !== undefined) {
-                                try {
-                                    finalOutputs = await aggregator(chunks);
+                            finally {
+                                if (!finished) {
+                                    await currentRunTree?.end(undefined, "Cancelled");
                                 }
-                                catch (e) {
-                                    console.error(`[ERROR]: LangSmith aggregation failed: `, e);
+                                let finalOutputs;
+                                if (aggregator !== undefined) {
+                                    try {
+                                        finalOutputs = await aggregator(chunks);
+                                    }
+                                    catch (e) {
+                                        console.error(`[ERROR]: LangSmith aggregation failed: `, e);
+                                        finalOutputs = chunks;
+                                    }
+                                }
+                                else {
                                     finalOutputs = chunks;
                                 }
+                                if (typeof finalOutputs === "object" &&
+                                    !Array.isArray(finalOutputs)) {
+                                    await currentRunTree?.end(finalOutputs);
+                                }
+                                else {
+                                    await currentRunTree?.end({ outputs: finalOutputs });
+                                }
+                                await postRunPromise;
+                                await currentRunTree?.patchRun();
                             }
-                            else {
-                                finalOutputs = chunks;
-                            }
-                            if (typeof finalOutputs === "object" &&
-                                !Array.isArray(finalOutputs)) {
-                                await currentRunTree?.end(finalOutputs);
-                            }
-                            else {
-                                await currentRunTree?.end({ outputs: finalOutputs });
-                            }
-                            await currentRunTree?.patchRun();
                         }
                         return resolve(wrapOutputForTracing());
                     }
                     else {
-                        const outputs = isKVMap(rawOutput)
-                            ? rawOutput
-                            : { outputs: rawOutput };
-                        if (initialOutputs === currentRunTree?.outputs) {
-                            await currentRunTree?.end(outputs);
-                        }
-                        else {
-                            if (currentRunTree !== undefined) {
-                                currentRunTree.end_time = Date.now();
-                            }
+                        try {
+                            await currentRunTree?.end(isKVMap(rawOutput) ? rawOutput : { outputs: rawOutput });
+                            await postRunPromise;
+                            await currentRunTree?.patchRun();
                         }
-                        await currentRunTree?.patchRun();
-                        return resolve(rawOutput);
-                    }
-                }
-                catch (error) {
-                    if (initialError === currentRunTree?.error) {
-                        await currentRunTree?.end(initialOutputs, String(error));
-                    }
-                    else {
-                        if (currentRunTree !== undefined) {
-                            currentRunTree.end_time = Date.now();
+                        finally {
+                            // eslint-disable-next-line no-unsafe-finally
+                            return rawOutput;
                         }
                     }
+                }, 
+                // eslint-disable-next-line @typescript-eslint/no-explicit-any
+                async (error) => {
+                    await currentRunTree?.end(undefined, String(error));
+                    await postRunPromise;
                     await currentRunTree?.patchRun();
-                    reject(error);
-                }
+                    throw error;
+                })
+                    .then(resolve, reject);
+            });
+            if (typeof returnValue !== "object" || returnValue === null) {
+                return tracedPromise;
+            }
+            return new Proxy(returnValue, {
+                get(target, prop, receiver) {
+                    if (isPromiseMethod(prop)) {
+                        return tracedPromise[prop].bind(tracedPromise);
+                    }
+                    return Reflect.get(target, prop, receiver);
+                },
             });
         });
     };

package/dist/utils/async_caller.cjs

@@ -7,13 +7,13 @@ exports.AsyncCaller = void 0;
 const p_retry_1 = __importDefault(require("p-retry"));
 const p_queue_1 = __importDefault(require("p-queue"));
 const STATUS_NO_RETRY = [
-    400,
-    401,
-    403,
-    404,
-    405,
-    406,
-    407,
+    400, // Bad Request
+    401, // Unauthorized
+    403, // Forbidden
+    404, // Not Found
+    405, // Method Not Allowed
+    406, // Not Acceptable
+    407, // Proxy Authentication Required
     408, // Request Timeout
 ];
 const STATUS_IGNORE = [
@@ -60,8 +60,16 @@ class AsyncCaller {
         });
         this.maxConcurrency = params.maxConcurrency ?? Infinity;
         this.maxRetries = params.maxRetries ?? 6;
-        const PQueue = "default" in p_queue_1.default ? p_queue_1.default.default : p_queue_1.default;
-        this.queue = new PQueue({ concurrency: this.maxConcurrency });
+        if ("default" in p_queue_1.default) {
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            this.queue = new p_queue_1.default.default({
+                concurrency: this.maxConcurrency,
+            });
+        }
+        else {
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            this.queue = new p_queue_1.default({ concurrency: this.maxConcurrency });
+        }
         this.onFailedResponseHook = params?.onFailedResponseHook;
     }
     // eslint-disable-next-line @typescript-eslint/no-explicit-any

package/dist/utils/async_caller.js

@@ -1,13 +1,13 @@
 import pRetry from "p-retry";
 import PQueueMod from "p-queue";
 const STATUS_NO_RETRY = [
-    400,
-    401,
-    403,
-    404,
-    405,
-    406,
-    407,
+    400, // Bad Request
+    401, // Unauthorized
+    403, // Forbidden
+    404, // Not Found
+    405, // Method Not Allowed
+    406, // Not Acceptable
+    407, // Proxy Authentication Required
     408, // Request Timeout
 ];
 const STATUS_IGNORE = [
@@ -54,8 +54,16 @@ export class AsyncCaller {
         });
         this.maxConcurrency = params.maxConcurrency ?? Infinity;
         this.maxRetries = params.maxRetries ?? 6;
-        const PQueue = "default" in PQueueMod ? PQueueMod.default : PQueueMod;
-        this.queue = new PQueue({ concurrency: this.maxConcurrency });
+        if ("default" in PQueueMod) {
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            this.queue = new PQueueMod.default({
+                concurrency: this.maxConcurrency,
+            });
+        }
+        else {
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            this.queue = new PQueueMod({ concurrency: this.maxConcurrency });
+        }
         this.onFailedResponseHook = params?.onFailedResponseHook;
     }
     // eslint-disable-next-line @typescript-eslint/no-explicit-any

package/dist/wrappers/openai.cjs

@@ -2,8 +2,11 @@
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.wrapSDK = exports.wrapOpenAI = void 0;
 const traceable_js_1 = require("../traceable.cjs");
-function _combineChatCompletionChoices(choices) {
+function _combineChatCompletionChoices(choices
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+) {
     const reversedChoices = choices.slice().reverse();
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
     const message = {
         role: "assistant",
         content: "",
@@ -73,26 +76,43 @@ function _combineChatCompletionChoices(choices) {
         message: message,
     };
 }
-async function extractLangSmithExtraAndCall(openAIMethod, args, defaultRunConfig) {
-    if (args[1]?.langsmithExtra !== undefined) {
-        const { langsmithExtra, ...openAIOptions } = args[1];
-        const wrappedMethod = (0, traceable_js_1.traceable)(openAIMethod, {
-            ...defaultRunConfig,
-            ...langsmithExtra,
-        });
-        const finalArgs = [args[0]];
-        if (args.length > 2) {
-            finalArgs.push(openAIOptions);
-            finalArgs.push(args.slice(2));
+const chatAggregator = (chunks) => {
+    if (!chunks || chunks.length === 0) {
+        return { choices: [{ message: { role: "assistant", content: "" } }] };
+    }
+    const choicesByIndex = {};
+    for (const chunk of chunks) {
+        for (const choice of chunk.choices) {
+            if (choicesByIndex[choice.index] === undefined) {
+                choicesByIndex[choice.index] = [];
+            }
+            choicesByIndex[choice.index].push(choice);
         }
-        else if (Object.keys(openAIOptions).length !== 0) {
-            finalArgs.push(openAIOptions);
+    }
+    const aggregatedOutput = chunks[chunks.length - 1];
+    aggregatedOutput.choices = Object.values(choicesByIndex).map((choices) => _combineChatCompletionChoices(choices));
+    return aggregatedOutput;
+};
+const textAggregator = (allChunks
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+) => {
+    if (allChunks.length === 0) {
+        return { choices: [{ text: "" }] };
+    }
+    const allContent = [];
+    for (const chunk of allChunks) {
+        const content = chunk.choices[0].text;
+        if (content != null) {
+            allContent.push(content);
         }
-        return wrappedMethod(...finalArgs);
     }
-    const wrappedMethod = (0, traceable_js_1.traceable)(openAIMethod, defaultRunConfig);
-    return wrappedMethod(...args);
-}
+    const content = allContent.join("");
+    const aggregatedOutput = allChunks[allChunks.length - 1];
+    aggregatedOutput.choices = [
+        { ...aggregatedOutput.choices[0], text: content },
+    ];
+    return aggregatedOutput;
+};
 /**
  * Wraps an OpenAI client's completion methods, enabling automatic LangSmith
  * tracing. Method signatures are unchanged, with the exception that you can pass
@@ -118,61 +138,20 @@ async function extractLangSmithExtraAndCall(openAIMethod, args, defaultRunConfig
  * ```
  */
 const wrapOpenAI = (openai, options) => {
-    const originalChatCompletionsFn = openai.chat.completions.create.bind(openai.chat.completions);
-    openai.chat.completions.create = async (...args) => {
-        const aggregator = (chunks) => {
-            if (!chunks || chunks.length === 0) {
-                return { choices: [{ message: { role: "assistant", content: "" } }] };
-            }
-            const choicesByIndex = {};
-            for (const chunk of chunks) {
-                for (const choice of chunk.choices) {
-                    if (choicesByIndex[choice.index] === undefined) {
-                        choicesByIndex[choice.index] = [];
-                    }
-                    choicesByIndex[choice.index].push(choice);
-                }
-            }
-            const aggregatedOutput = chunks[chunks.length - 1];
-            aggregatedOutput.choices = Object.values(choicesByIndex).map((choices) => _combineChatCompletionChoices(choices));
-            return aggregatedOutput;
-        };
-        const defaultRunConfig = {
-            name: "ChatOpenAI",
-            run_type: "llm",
-            aggregator,
-            ...options,
-        };
-        return extractLangSmithExtraAndCall(originalChatCompletionsFn, args, defaultRunConfig);
-    };
-    const originalCompletionsFn = openai.completions.create.bind(openai.chat.completions);
-    openai.completions.create = async (...args) => {
-        const aggregator = (allChunks) => {
-            if (allChunks.length === 0) {
-                return { choices: [{ text: "" }] };
-            }
-            const allContent = [];
-            for (const chunk of allChunks) {
-                const content = chunk.choices[0].text;
-                if (content != null) {
-                    allContent.push(content);
-                }
-            }
-            const content = allContent.join("");
-            const aggregatedOutput = allChunks[allChunks.length - 1];
-            aggregatedOutput.choices = [
-                { ...aggregatedOutput.choices[0], text: content },
-            ];
-            return aggregatedOutput;
-        };
-        const defaultRunConfig = {
-            name: "OpenAI",
-            run_type: "llm",
-            aggregator,
-            ...options,
-        };
-        return extractLangSmithExtraAndCall(originalCompletionsFn, args, defaultRunConfig);
-    };
+    openai.chat.completions.create = (0, traceable_js_1.traceable)(openai.chat.completions.create.bind(openai.chat.completions), {
+        name: "ChatOpenAI",
+        run_type: "llm",
+        aggregator: chatAggregator,
+        argsConfigPath: [1, "langsmithExtra"],
+        ...options,
+    });
+    openai.completions.create = (0, traceable_js_1.traceable)(openai.completions.create.bind(openai.completions), {
+        name: "OpenAI",
+        run_type: "llm",
+        aggregator: textAggregator,
+        argsConfigPath: [1, "langsmithExtra"],
+        ...options,
+    });
     return openai;
 };
 exports.wrapOpenAI = wrapOpenAI;
@@ -181,7 +160,7 @@ const _wrapClient = (sdk, runName, options) => {
         get(target, propKey, receiver) {
             const originalValue = target[propKey];
             if (typeof originalValue === "function") {
-                return (0, traceable_js_1.traceable)(originalValue.bind(target), Object.assign({ name: [runName, propKey.toString()].join("."), run_type: "llm" }, options?.client));
+                return (0, traceable_js_1.traceable)(originalValue.bind(target), Object.assign({ name: [runName, propKey.toString()].join("."), run_type: "llm" }, options));
             }
             else if (originalValue != null &&
                 !Array.isArray(originalValue) &&

package/dist/wrappers/openai.d.ts

@@ -1,4 +1,5 @@
-import type { OpenAI } from "openai";
+import { OpenAI } from "openai";
+import type { APIPromise } from "openai/core";
 import type { Client, RunTreeConfig } from "../index.js";
 import { type RunnableConfigLike } from "../run_trees.js";
 import { type RunTreeLike } from "../traceable.js";
@@ -12,31 +13,29 @@ type OpenAIType = {
         create: (...args: any[]) => any;
     };
 };
-type PatchedOpenAIClient<T extends OpenAIType> = {
-    [P in keyof T]: T[P];
-} & {
-    chat: {
-        completions: {
+type PatchedOpenAIClient<T extends OpenAIType> = T & {
+    chat: T["chat"] & {
+        completions: T["chat"]["completions"] & {
             create: {
                 (arg: OpenAI.ChatCompletionCreateParamsStreaming, arg2?: OpenAI.RequestOptions & {
                     langsmithExtra?: RunnableConfigLike | RunTreeLike;
-                }): Promise<AsyncGenerator<OpenAI.ChatCompletionChunk>>;
+                }): APIPromise<AsyncGenerator<OpenAI.ChatCompletionChunk>>;
             } & {
                 (arg: OpenAI.ChatCompletionCreateParamsNonStreaming, arg2?: OpenAI.RequestOptions & {
                     langsmithExtra?: RunnableConfigLike | RunTreeLike;
-                }): Promise<OpenAI.ChatCompletionChunk>;
+                }): APIPromise<OpenAI.ChatCompletionChunk>;
             };
         };
     };
-    completions: {
+    completions: T["completions"] & {
         create: {
             (arg: OpenAI.CompletionCreateParamsStreaming, arg2?: OpenAI.RequestOptions & {
                 langsmithExtra?: RunnableConfigLike | RunTreeLike;
-            }): Promise<AsyncGenerator<OpenAI.Completion>>;
+            }): APIPromise<AsyncGenerator<OpenAI.Completion>>;
         } & {
             (arg: OpenAI.CompletionCreateParamsNonStreaming, arg2?: OpenAI.RequestOptions & {
                 langsmithExtra?: RunnableConfigLike | RunTreeLike;
-            }): Promise<OpenAI.Completion>;
+            }): APIPromise<OpenAI.Completion>;
         };
     };
 };

package/dist/wrappers/openai.js

@@ -1,6 +1,9 @@
 import { traceable } from "../traceable.js";
-function _combineChatCompletionChoices(choices) {
+function _combineChatCompletionChoices(choices
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+) {
     const reversedChoices = choices.slice().reverse();
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
     const message = {
         role: "assistant",
         content: "",
@@ -70,26 +73,43 @@ function _combineChatCompletionChoices(choices) {
         message: message,
     };
 }
-async function extractLangSmithExtraAndCall(openAIMethod, args, defaultRunConfig) {
-    if (args[1]?.langsmithExtra !== undefined) {
-        const { langsmithExtra, ...openAIOptions } = args[1];
-        const wrappedMethod = traceable(openAIMethod, {
-            ...defaultRunConfig,
-            ...langsmithExtra,
-        });
-        const finalArgs = [args[0]];
-        if (args.length > 2) {
-            finalArgs.push(openAIOptions);
-            finalArgs.push(args.slice(2));
+const chatAggregator = (chunks) => {
+    if (!chunks || chunks.length === 0) {
+        return { choices: [{ message: { role: "assistant", content: "" } }] };
+    }
+    const choicesByIndex = {};
+    for (const chunk of chunks) {
+        for (const choice of chunk.choices) {
+            if (choicesByIndex[choice.index] === undefined) {
+                choicesByIndex[choice.index] = [];
+            }
+            choicesByIndex[choice.index].push(choice);
         }
-        else if (Object.keys(openAIOptions).length !== 0) {
-            finalArgs.push(openAIOptions);
+    }
+    const aggregatedOutput = chunks[chunks.length - 1];
+    aggregatedOutput.choices = Object.values(choicesByIndex).map((choices) => _combineChatCompletionChoices(choices));
+    return aggregatedOutput;
+};
+const textAggregator = (allChunks
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+) => {
+    if (allChunks.length === 0) {
+        return { choices: [{ text: "" }] };
+    }
+    const allContent = [];
+    for (const chunk of allChunks) {
+        const content = chunk.choices[0].text;
+        if (content != null) {
+            allContent.push(content);
         }
-        return wrappedMethod(...finalArgs);
     }
-    const wrappedMethod = traceable(openAIMethod, defaultRunConfig);
-    return wrappedMethod(...args);
-}
+    const content = allContent.join("");
+    const aggregatedOutput = allChunks[allChunks.length - 1];
+    aggregatedOutput.choices = [
+        { ...aggregatedOutput.choices[0], text: content },
+    ];
+    return aggregatedOutput;
+};
 /**
  * Wraps an OpenAI client's completion methods, enabling automatic LangSmith
  * tracing. Method signatures are unchanged, with the exception that you can pass
@@ -115,61 +135,20 @@ async function extractLangSmithExtraAndCall(openAIMethod, args, defaultRunConfig
  * ```
  */
 export const wrapOpenAI = (openai, options) => {
-    const originalChatCompletionsFn = openai.chat.completions.create.bind(openai.chat.completions);
-    openai.chat.completions.create = async (...args) => {
-        const aggregator = (chunks) => {
-            if (!chunks || chunks.length === 0) {
-                return { choices: [{ message: { role: "assistant", content: "" } }] };
-            }
-            const choicesByIndex = {};
-            for (const chunk of chunks) {
-                for (const choice of chunk.choices) {
-                    if (choicesByIndex[choice.index] === undefined) {
-                        choicesByIndex[choice.index] = [];
-                    }
-                    choicesByIndex[choice.index].push(choice);
-                }
-            }
-            const aggregatedOutput = chunks[chunks.length - 1];
-            aggregatedOutput.choices = Object.values(choicesByIndex).map((choices) => _combineChatCompletionChoices(choices));
-            return aggregatedOutput;
-        };
-        const defaultRunConfig = {
-            name: "ChatOpenAI",
-            run_type: "llm",
-            aggregator,
-            ...options,
-        };
-        return extractLangSmithExtraAndCall(originalChatCompletionsFn, args, defaultRunConfig);
-    };
-    const originalCompletionsFn = openai.completions.create.bind(openai.chat.completions);
-    openai.completions.create = async (...args) => {
-        const aggregator = (allChunks) => {
-            if (allChunks.length === 0) {
-                return { choices: [{ text: "" }] };
-            }
-            const allContent = [];
-            for (const chunk of allChunks) {
-                const content = chunk.choices[0].text;
-                if (content != null) {
-                    allContent.push(content);
-                }
-            }
-            const content = allContent.join("");
-            const aggregatedOutput = allChunks[allChunks.length - 1];
-            aggregatedOutput.choices = [
-                { ...aggregatedOutput.choices[0], text: content },
-            ];
-            return aggregatedOutput;
-        };
-        const defaultRunConfig = {
-            name: "OpenAI",
-            run_type: "llm",
-            aggregator,
-            ...options,
-        };
-        return extractLangSmithExtraAndCall(originalCompletionsFn, args, defaultRunConfig);
-    };
+    openai.chat.completions.create = traceable(openai.chat.completions.create.bind(openai.chat.completions), {
+        name: "ChatOpenAI",
+        run_type: "llm",
+        aggregator: chatAggregator,
+        argsConfigPath: [1, "langsmithExtra"],
+        ...options,
+    });
+    openai.completions.create = traceable(openai.completions.create.bind(openai.completions), {
+        name: "OpenAI",
+        run_type: "llm",
+        aggregator: textAggregator,
+        argsConfigPath: [1, "langsmithExtra"],
+        ...options,
+    });
     return openai;
 };
 const _wrapClient = (sdk, runName, options) => {
@@ -177,7 +156,7 @@ const _wrapClient = (sdk, runName, options) => {
         get(target, propKey, receiver) {
             const originalValue = target[propKey];
             if (typeof originalValue === "function") {
-                return traceable(originalValue.bind(target), Object.assign({ name: [runName, propKey.toString()].join("."), run_type: "llm" }, options?.client));
+                return traceable(originalValue.bind(target), Object.assign({ name: [runName, propKey.toString()].join("."), run_type: "llm" }, options));
             }
             else if (originalValue != null &&
                 !Array.isArray(originalValue) &&

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "langsmith",
-  "version": "0.1.20",
+  "version": "0.1.21",
   "description": "Client library to connect to the LangSmith LLM Tracing and Evaluation Platform.",
   "packageManager": "yarn@1.22.19",
   "files": [
@@ -47,8 +47,8 @@
     "check-version": "node scripts/check-version.js",
     "check-npm-version": "node scripts/check-npm-version.js",
     "clean": "rm -rf dist/ && node scripts/create-entrypoints.js clean",
-    "build:esm": "tsc --outDir dist/ && rm -rf dist/tests dist/**/tests",
-    "build:cjs": "tsc --outDir dist-cjs/ -p tsconfig.cjs.json && node scripts/move-cjs-to-dist.js && rm -r dist-cjs",
+    "build:esm": "rm -f src/package.json && tsc --outDir dist/ && rm -rf dist/tests dist/**/tests",
+    "build:cjs": "echo '{}' > src/package.json && tsc --outDir dist-cjs/ -p tsconfig.cjs.json && node scripts/move-cjs-to-dist.js && rm -r dist-cjs src/package.json",
     "test": "cross-env NODE_OPTIONS=--experimental-vm-modules jest --passWithNoTests --testPathIgnorePatterns='\\.int\\.test.[tj]s' --testTimeout 30000",
     "test:integration": "cross-env NODE_OPTIONS=--experimental-vm-modules jest --testPathPattern=\\.int\\.test.ts --testTimeout 100000",
     "test:single": "NODE_OPTIONS=--experimental-vm-modules yarn run jest --config jest.config.cjs --testTimeout 100000",
@@ -104,7 +104,7 @@
     "prettier": "^2.8.8",
     "ts-jest": "^29.1.0",
     "ts-node": "^10.9.1",
-    "typescript": "^5.0.4"
+    "typescript": "^5.4.5"
   },
   "peerDependencies": {
     "openai": "*"