Skip to content

Commit

Permalink
Cleanup
Browse files Browse the repository at this point in the history
  • Loading branch information
@dqbd
dqbd committed May 16, 2024
1 parent 24a99c8 commit 3886649
Show file tree
Hide file tree
Showing 3 changed files with 133 additions and 127 deletions.
125 changes: 33 additions & 92 deletions js/src/singletons/traceable.ts
View file
Original file line number Diff line number Diff line change
@@ -1,24 +1,42 @@
import { RunTree, RunnableConfigLike } from "../run_trees.js";
import { RunTree } from "../run_trees.js";
import { TraceableFunction } from "./types.js";

interface AsyncStorageLike {
interface AsyncLocalStorageInterface {
getStore: () => RunTree | undefined;

run: (context: RunTree | undefined, fn: () => void) => void;
}

export const TraceableLocalStorageContext = (() => {
let storage: AsyncStorageLike;
class MockAsyncLocalStorage implements AsyncLocalStorageInterface {
getStore() {
return undefined;
}

run(_: RunTree | undefined, callback: () => void): void {
return callback();
}
}

class AsyncLocalStorageProvider {
private asyncLocalStorage: AsyncLocalStorageInterface =
new MockAsyncLocalStorage();

return {
register: (value: AsyncStorageLike) => {
storage ??= value;
return storage;
},
get storage() {
return storage;
},
};
})();
private hasBeenInitialized = false;

getInstance(): AsyncLocalStorageInterface {
return this.asyncLocalStorage;
}

initializeGlobalInstance(instance: AsyncLocalStorageInterface) {
if (!this.hasBeenInitialized) {
this.hasBeenInitialized = true;
this.asyncLocalStorage = instance;
}
}
}

export const AsyncLocalStorageProviderSingleton =
new AsyncLocalStorageProvider();

/**
* Return the current run tree from within a traceable-wrapped function.
Expand All @@ -27,11 +45,7 @@ export const TraceableLocalStorageContext = (() => {
* @returns The run tree for the given context.
*/
export const getCurrentRunTree = () => {
if (!TraceableLocalStorageContext.storage) {
throw new Error("Could not find the traceable storage context");
}

const runTree = TraceableLocalStorageContext.storage.getStore();
const runTree = AsyncLocalStorageProviderSingleton.getInstance().getStore();
if (runTree === undefined) {
throw new Error(
[
Expand All @@ -47,79 +61,6 @@ export const getCurrentRunTree = () => {

export const ROOT = Symbol.for("langsmith:traceable:root");

type SmartPromise<T> = T extends AsyncGenerator
? T
: T extends Promise<unknown>
? T
: Promise<T>;

type WrapArgReturnPair<Pair> = Pair extends [
// eslint-disable-next-line @typescript-eslint/no-explicit-any
infer Args extends any[],
infer Return
]
? Args extends [RunTree, ...infer RestArgs]
? {
(
runTree: RunTree | typeof ROOT,
...args: RestArgs
): SmartPromise<Return>;
(config: RunnableConfigLike, ...args: RestArgs): SmartPromise<Return>;
}
: {
(...args: Args): SmartPromise<Return>;
(runTree: RunTree, ...rest: Args): SmartPromise<Return>;
(config: RunnableConfigLike, ...args: Args): SmartPromise<Return>;
}
: never;

// eslint-disable-next-line @typescript-eslint/no-explicit-any
type UnionToIntersection<U> = (U extends any ? (x: U) => void : never) extends (
x: infer I
) => void
? I
: never;

// eslint-disable-next-line @typescript-eslint/no-explicit-any
export type TraceableFunction<Func extends (...args: any[]) => any> =
// function overloads are represented as intersections rather than unions
// matches the behavior introduced in https://github.com/microsoft/TypeScript/pull/54448
Func extends {
(...args: infer A1): infer R1;
(...args: infer A2): infer R2;
(...args: infer A3): infer R3;
(...args: infer A4): infer R4;
(...args: infer A5): infer R5;
}
? UnionToIntersection<
WrapArgReturnPair<[A1, R1] | [A2, R2] | [A3, R3] | [A4, R4] | [A5, R5]>
>
: Func extends {
(...args: infer A1): infer R1;
(...args: infer A2): infer R2;
(...args: infer A3): infer R3;
(...args: infer A4): infer R4;
}
? UnionToIntersection<
WrapArgReturnPair<[A1, R1] | [A2, R2] | [A3, R3] | [A4, R4]>
>
: Func extends {
(...args: infer A1): infer R1;
(...args: infer A2): infer R2;
(...args: infer A3): infer R3;
}
? UnionToIntersection<WrapArgReturnPair<[A1, R1] | [A2, R2] | [A3, R3]>>
: Func extends {
(...args: infer A1): infer R1;
(...args: infer A2): infer R2;
}
? UnionToIntersection<WrapArgReturnPair<[A1, R1] | [A2, R2]>>
: Func extends {
(...args: infer A1): infer R1;
}
? UnionToIntersection<WrapArgReturnPair<[A1, R1]>>
: never;

export function isTraceableFunction(
x: unknown
// eslint-disable-next-line @typescript-eslint/no-explicit-any
Expand Down
75 changes: 75 additions & 0 deletions js/src/singletons/types.ts
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
import { RunTree, RunnableConfigLike } from "../run_trees.js";
import { ROOT } from "./traceable.js";

type SmartPromise<T> = T extends AsyncGenerator
? T
: T extends Promise<unknown>
? T
: Promise<T>;
type WrapArgReturnPair<Pair> = Pair extends [
// eslint-disable-next-line @typescript-eslint/no-explicit-any
infer Args extends any[],
infer Return
]
? Args extends [RunTree, ...infer RestArgs]
? {
(
runTree: RunTree | typeof ROOT,
...args: RestArgs
): SmartPromise<Return>;
(config: RunnableConfigLike, ...args: RestArgs): SmartPromise<Return>;
}
: {
(...args: Args): SmartPromise<Return>;
(runTree: RunTree, ...rest: Args): SmartPromise<Return>;
(config: RunnableConfigLike, ...args: Args): SmartPromise<Return>;
}
: never;
// eslint-disable-next-line @typescript-eslint/no-explicit-any
type UnionToIntersection<U> = (U extends any ? (x: U) => void : never) extends (
x: infer I
) => void
? I
: never;
// eslint-disable-next-line @typescript-eslint/no-explicit-any

export type TraceableFunction<Func extends (...args: any[]) => any> =
// function overloads are represented as intersections rather than unions
// matches the behavior introduced in https://github.com/microsoft/TypeScript/pull/54448
Func extends {
(...args: infer A1): infer R1;
(...args: infer A2): infer R2;
(...args: infer A3): infer R3;
(...args: infer A4): infer R4;
(...args: infer A5): infer R5;
}
? UnionToIntersection<
WrapArgReturnPair<[A1, R1] | [A2, R2] | [A3, R3] | [A4, R4] | [A5, R5]>
>
: Func extends {
(...args: infer A1): infer R1;
(...args: infer A2): infer R2;
(...args: infer A3): infer R3;
(...args: infer A4): infer R4;
}
? UnionToIntersection<
WrapArgReturnPair<[A1, R1] | [A2, R2] | [A3, R3] | [A4, R4]>
>
: Func extends {
(...args: infer A1): infer R1;
(...args: infer A2): infer R2;
(...args: infer A3): infer R3;
}
? UnionToIntersection<WrapArgReturnPair<[A1, R1] | [A2, R2] | [A3, R3]>>
: Func extends {
(...args: infer A1): infer R1;
(...args: infer A2): infer R2;
}
? UnionToIntersection<WrapArgReturnPair<[A1, R1] | [A2, R2]>>
: Func extends {
(...args: infer A1): infer R1;
}
? UnionToIntersection<WrapArgReturnPair<[A1, R1]>>
: never;

export type RunTreeLike = RunTree;
60 changes: 25 additions & 35 deletions js/src/traceable.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,9 +11,13 @@ import { InvocationParamsSchema, KVMap } from "./schemas.js";
import { isTracingEnabled } from "./env.js";
import {
ROOT,
TraceableFunction,
TraceableLocalStorageContext,
AsyncLocalStorageProviderSingleton,
} from "./singletons/traceable.js";
import { TraceableFunction } from "./singletons/types.js";

AsyncLocalStorageProviderSingleton.initializeGlobalInstance(
new AsyncLocalStorage<RunTree | undefined>()
);

function isPromiseMethod(
x: string | symbol
Expand All @@ -24,24 +28,35 @@ function isPromiseMethod(
return false;
}

const asyncLocalStorage = TraceableLocalStorageContext.register(
new AsyncLocalStorage<RunTree | undefined>()
);
function isKVMap(x: unknown): x is Record<string, unknown> {
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)
);
}

const isAsyncIterable = (x: unknown): x is AsyncIterable<unknown> =>
x != null &&
typeof x === "object" &&
// eslint-disable-next-line @typescript-eslint/no-explicit-any
typeof (x as any)[Symbol.asyncIterator] === "function";

const GeneratorFunction = function* () {}.constructor;

const isIteratorLike = (x: unknown): x is Iterator<unknown> =>
x != null &&
typeof x === "object" &&
"next" in x &&
typeof x.next === "function";

const GeneratorFunction = function* () {}.constructor;

const isGenerator = (x: unknown): x is Generator =>
// eslint-disable-next-line no-instanceof/no-instanceof
x != null && typeof x === "function" && x instanceof GeneratorFunction;
Expand Down Expand Up @@ -379,6 +394,8 @@ export function traceable<Func extends (...args: any[]) => any>(
};
}

const asyncLocalStorage = AsyncLocalStorageProviderSingleton.getInstance();

// TODO: deal with possible nested promises and async iterables
const processedArgs = args as unknown as Inputs;
for (let i = 0; i < processedArgs.length; i++) {
Expand Down Expand Up @@ -608,36 +625,9 @@ export function traceable<Func extends (...args: any[]) => any>(
}

export {
type TraceableFunction,
getCurrentRunTree,
isTraceableFunction,
ROOT,
} from "./singletons/traceable.js";
export type RunTreeLike = RunTree;

function isKVMap(x: unknown): x is Record<string, unknown> {
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)
);
}

export function wrapFunctionAndEnsureTraceable<
Func extends (...args: any[]) => any
>(target: Func, options: Partial<RunTreeConfig>, name = "target") {
if (typeof target === "function") {
return traceable<Func>(target, {
...options,
name,
});
}
throw new Error("Target must be runnable function");
}
export type { RunTreeLike, TraceableFunction } from "./singletons/types.js";

0 comments on commit 3886649

Please sign in to comment.