langsmith 0.1.20 → 0.1.22

package/dist/traceable.cjs

@@ -1,18 +1,35 @@
 "use strict";
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.isTraceableFunction = exports.getCurrentRunTree = exports.traceable = void 0;
+exports.wrapFunctionAndEnsureTraceable = exports.isTraceableFunction = exports.getCurrentRunTree = exports.traceable = void 0;
 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" &&
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     typeof x[Symbol.asyncIterator] === "function";
-const getTracingRunTree = (runTree) => {
-    const tracingEnabled = (0, env_js_1.getEnvironmentVariable)("LANGSMITH_TRACING_V2") === "true" ||
-        (0, env_js_1.getEnvironmentVariable)("LANGCHAIN_TRACING_V2") === "true";
-    if (!tracingEnabled) {
+const tracingIsEnabled = (tracingEnabled) => {
+    if (tracingEnabled !== undefined) {
+        return tracingEnabled;
+    }
+    const envVars = [
+        "LANGSMITH_TRACING_V2",
+        "LANGCHAIN_TRACING_V2",
+        "LANGSMITH_TRACING",
+        "LANGCHAIN_TRACING",
+    ];
+    return Boolean(envVars.find((envVar) => (0, env_js_1.getEnvironmentVariable)(envVar) === "true"));
+};
+const getTracingRunTree = (runTree, tracingEnabled) => {
+    const tracingEnabled_ = tracingIsEnabled(tracingEnabled);
+    if (!tracingEnabled_) {
         return undefined;
     }
     return runTree;
@@ -33,14 +50,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, tracingEnabled, ...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,14 +109,14 @@ function traceable(wrappedFunc, config) {
             rawInputs = args.slice(1);
         }
         else if (previousRunTree !== undefined) {
-            currentRunTree = await previousRunTree.createChild(ensuredConfig);
+            currentRunTree = previousRunTree.createChild(ensuredConfig);
             rawInputs = args;
         }
         else {
             currentRunTree = new run_trees_js_1.RunTree(ensuredConfig);
             rawInputs = args;
         }
-        currentRunTree = getTracingRunTree(currentRunTree);
+        currentRunTree = getTracingRunTree(currentRunTree, tracingEnabled);
         let inputs;
         const firstInput = rawInputs[0];
         if (firstInput == null) {
@@ -76,74 +134,181 @@ 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 });
+                        }
+                        const onEnd = config?.on_end;
+                        if (onEnd) {
+                            if (!currentRunTree) {
+                                console.warn("Can not call 'on_end' if currentRunTree is undefined");
+                            }
+                            else {
+                                onEnd(currentRunTree);
+                            }
+                        }
+                        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 });
+                                }
+                                const onEnd = config?.on_end;
+                                if (onEnd) {
+                                    if (!currentRunTree) {
+                                        console.warn("Can not call 'on_end' if currentRunTree is undefined");
+                                    }
+                                    else {
+                                        onEnd(currentRunTree);
+                                    }
+                                }
+                                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 });
+                            const onEnd = config?.on_end;
+                            if (onEnd) {
+                                if (!currentRunTree) {
+                                    console.warn("Can not call 'on_end' if currentRunTree is undefined");
+                                }
+                                else {
+                                    onEnd(currentRunTree);
+                                }
                             }
+                            await postRunPromise;
+                            await currentRunTree?.patchRun();
+                        }
+                        finally {
+                            // eslint-disable-next-line no-unsafe-finally
+                            return rawOutput;
                         }
-                        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();
+                }, 
+                // eslint-disable-next-line @typescript-eslint/no-explicit-any
+                async (error) => {
+                    await currentRunTree?.end(undefined, String(error));
+                    const onEnd = config?.on_end;
+                    if (onEnd) {
+                        if (!currentRunTree) {
+                            console.warn("Can not call 'on_end' if currentRunTree is undefined");
+                        }
+                        else {
+                            onEnd(currentRunTree);
                         }
                     }
+                    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);
+                },
             });
         });
     };
@@ -184,3 +349,13 @@ function isKVMap(x) {
         // eslint-disable-next-line no-instanceof/no-instanceof
         !(x instanceof Date));
 }
+function wrapFunctionAndEnsureTraceable(target, options, name = "target") {
+    if (typeof target === "function") {
+        return traceable(target, {
+            ...options,
+            name,
+        });
+    }
+    throw new Error("Target must be runnable function");
+}
+exports.wrapFunctionAndEnsureTraceable = wrapFunctionAndEnsureTraceable;

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,8 @@ 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];
+    tracingEnabled?: boolean;
 }): TraceableFunction<Func>;
 /**
  * Return the current run tree from within a traceable-wrapped function.
@@ -55,4 +58,5 @@ export declare function traceable<Func extends (...args: any[]) => any>(wrappedF
  */
 export declare function getCurrentRunTree(): RunTree;
 export declare function isTraceableFunction(x: unknown): x is TraceableFunction<any>;
+export declare function wrapFunctionAndEnsureTraceable<Func extends (...args: any[]) => any>(target: Func, options: Partial<RunTreeConfig>, name?: string): TraceableFunction<Func>;
 export {};