@mondaydotcomorg/atp-compiler 0.17.14

package/src/transformer/promise-transformer.ts

@@ -0,0 +1,194 @@
+import * as t from '@babel/types';
+import { generateUniqueId } from '../runtime/context.js';
+import { BatchParallelDetector } from './batch-detector.js';
+import type { BatchCallInfo } from '../types.js';
+
+export class PromiseTransformer {
+	private transformCount = 0;
+	private batchDetector: BatchParallelDetector;
+	private enableBatchParallel: boolean;
+
+	constructor(enableBatchParallel = true) {
+		this.batchDetector = new BatchParallelDetector();
+		this.enableBatchParallel = enableBatchParallel;
+	}
+
+	transformPromiseAll(path: any): boolean {
+		const node = path.node as t.CallExpression;
+
+		if (!this.isPromiseAll(node)) {
+			return false;
+		}
+
+		const arrayArg = node.arguments[0];
+
+		if (this.enableBatchParallel && this.batchDetector.canBatch(node)) {
+			return this.transformToBatchParallel(path, node);
+		}
+
+		if (t.isArrayExpression(arrayArg)) {
+			return this.transformToSequential(path, node);
+		}
+
+		return false;
+	}
+
+	transformPromiseAllSettled(path: any): boolean {
+		const node = path.node as t.CallExpression;
+
+		if (!this.isPromiseAllSettled(node)) {
+			return false;
+		}
+
+		const arrayArg = node.arguments[0];
+
+		if (t.isArrayExpression(arrayArg)) {
+			const parallelId = generateUniqueId('allSettled');
+
+			const runtimeCall = t.awaitExpression(
+				t.callExpression(
+					t.memberExpression(t.identifier('__runtime'), t.identifier('resumablePromiseAllSettled')),
+					[arrayArg, t.stringLiteral(parallelId)]
+				)
+			);
+
+			path.replaceWith(runtimeCall);
+			this.transformCount++;
+			return true;
+		}
+
+		return false;
+	}
+
+	private transformToBatchParallel(path: any, node: t.CallExpression): boolean {
+		const arrayArg = node.arguments[0];
+		if (!t.isArrayExpression(arrayArg)) {
+			return false;
+		}
+
+		const batchId = generateUniqueId('batch');
+
+		const batchCallsArray = t.arrayExpression(
+			arrayArg.elements.map((el) => {
+				if (!el || t.isSpreadElement(el)) {
+					return t.nullLiteral();
+				}
+
+				let callNode: t.Node = el;
+				if (t.isAwaitExpression(callNode)) {
+					callNode = callNode.argument;
+				}
+
+				if (!t.isCallExpression(callNode) || !t.isMemberExpression(callNode.callee)) {
+					return t.nullLiteral();
+				}
+
+				const callInfo = this.batchDetector.extractCallInfo(callNode);
+				if (!callInfo) {
+					return t.nullLiteral();
+				}
+
+				const payloadArg = callNode.arguments[0];
+
+				return t.objectExpression([
+					t.objectProperty(t.identifier('type'), t.stringLiteral(callInfo.type)),
+					t.objectProperty(t.identifier('operation'), t.stringLiteral(callInfo.operation)),
+					t.objectProperty(
+						t.identifier('payload'),
+						payloadArg && t.isExpression(payloadArg) ? payloadArg : t.objectExpression([])
+					),
+				]);
+			})
+		);
+
+		const runtimeCall = t.awaitExpression(
+			t.callExpression(
+				t.memberExpression(t.identifier('__runtime'), t.identifier('batchParallel')),
+				[batchCallsArray, t.stringLiteral(batchId)]
+			)
+		);
+
+		path.replaceWith(runtimeCall);
+		this.transformCount++;
+		return true;
+	}
+
+	private transformToSequential(path: any, node: t.CallExpression): boolean {
+		const arrayArg = node.arguments[0];
+		if (!t.isArrayExpression(arrayArg)) {
+			return false;
+		}
+
+		const parallelId = generateUniqueId('parallel');
+
+		const runtimeCall = t.awaitExpression(
+			t.callExpression(
+				t.memberExpression(t.identifier('__runtime'), t.identifier('resumablePromiseAll')),
+				[arrayArg, t.stringLiteral(parallelId)]
+			)
+		);
+
+		path.replaceWith(runtimeCall);
+		this.transformCount++;
+		return true;
+	}
+
+	private payloadToExpression(payload: Record<string, unknown>): t.ObjectExpression {
+		const properties: t.ObjectProperty[] = [];
+
+		for (const [key, value] of Object.entries(payload)) {
+			properties.push(t.objectProperty(t.identifier(key), this.valueToExpression(value)));
+		}
+
+		return t.objectExpression(properties);
+	}
+
+	private valueToExpression(value: unknown): t.Expression {
+		if (typeof value === 'string') {
+			return t.stringLiteral(value);
+		}
+		if (typeof value === 'number') {
+			return t.numericLiteral(value);
+		}
+		if (typeof value === 'boolean') {
+			return t.booleanLiteral(value);
+		}
+		if (value === null) {
+			return t.nullLiteral();
+		}
+		if (Array.isArray(value)) {
+			return t.arrayExpression(value.map((v) => this.valueToExpression(v)));
+		}
+		if (typeof value === 'object') {
+			return this.payloadToExpression(value as Record<string, unknown>);
+		}
+
+		return t.identifier('undefined');
+	}
+
+	private isPromiseAll(node: t.CallExpression): boolean {
+		const callee = node.callee;
+		return (
+			t.isMemberExpression(callee) &&
+			t.isIdentifier(callee.object, { name: 'Promise' }) &&
+			t.isIdentifier(callee.property, { name: 'all' })
+		);
+	}
+
+	private isPromiseAllSettled(node: t.CallExpression): boolean {
+		const callee = node.callee;
+		return (
+			t.isMemberExpression(callee) &&
+			t.isIdentifier(callee.object, { name: 'Promise' }) &&
+			t.isIdentifier(callee.property, { name: 'allSettled' })
+		);
+	}
+
+	getTransformCount(): number {
+		return this.transformCount;
+	}
+
+	resetTransformCount(): void {
+		this.transformCount = 0;
+	}
+}

package/src/transformer/utils.ts

@@ -0,0 +1,147 @@
+import * as t from '@babel/types';
+import { PAUSABLE_CALL_PATTERNS } from '../types.js';
+
+export function isPausableCall(node: t.Node): boolean {
+	if (!t.isAwaitExpression(node)) {
+		return false;
+	}
+
+	const argument = node.argument;
+	if (!t.isCallExpression(argument)) {
+		return false;
+	}
+
+	return isPausableCallExpression(argument);
+}
+
+export function isPausableCallExpression(node: t.CallExpression): boolean {
+	const callee = node.callee;
+
+	if (!t.isMemberExpression(callee)) {
+		return false;
+	}
+
+	const fullPath = getMemberExpressionPath(callee);
+
+	return PAUSABLE_CALL_PATTERNS.some(
+		(pattern) => fullPath === `${pattern.namespace}.${pattern.method}`
+	);
+}
+
+export function getMemberExpressionPath(node: t.MemberExpression): string {
+	const parts: string[] = [];
+
+	let current: t.Node = node;
+	while (t.isMemberExpression(current)) {
+		if (t.isIdentifier(current.property)) {
+			parts.unshift(current.property.name);
+		}
+		current = current.object;
+	}
+
+	if (t.isIdentifier(current)) {
+		parts.unshift(current.name);
+	}
+
+	return parts.join('.');
+}
+
+export function containsAwait(node: t.Node): boolean {
+	let hasAwait = false;
+
+	const checkNode = (n: t.Node): void => {
+		if (t.isAwaitExpression(n)) {
+			hasAwait = true;
+			return;
+		}
+
+		if (hasAwait) return;
+
+		Object.keys(n).forEach((key) => {
+			const value = (n as any)[key];
+			if (Array.isArray(value)) {
+				value.forEach((item) => {
+					if (item && typeof item === 'object' && item.type) {
+						checkNode(item);
+					}
+				});
+			} else if (value && typeof value === 'object' && value.type) {
+				checkNode(value);
+			}
+		});
+	};
+
+	checkNode(node);
+	return hasAwait;
+}
+
+export function containsPausableCall(node: t.Node): boolean {
+	let hasPausable = false;
+
+	const checkNode = (n: t.Node): void => {
+		if (t.isAwaitExpression(n) && isPausableCall(n)) {
+			hasPausable = true;
+			return;
+		}
+
+		if (hasPausable) return;
+
+		Object.keys(n).forEach((key) => {
+			const value = (n as any)[key];
+			if (Array.isArray(value)) {
+				value.forEach((item) => {
+					if (item && typeof item === 'object' && item.type) {
+						checkNode(item);
+					}
+				});
+			} else if (value && typeof value === 'object' && value.type) {
+				checkNode(value);
+			}
+		});
+	};
+
+	checkNode(node);
+	return hasPausable;
+}
+
+export function isAsyncFunction(node: t.Node): boolean {
+	return (
+		(t.isFunctionDeclaration(node) ||
+			t.isFunctionExpression(node) ||
+			t.isArrowFunctionExpression(node)) &&
+		node.async === true
+	);
+}
+
+export function getNodeLocation(node: t.Node): { line: number; column: number } | undefined {
+	if (node.loc) {
+		return {
+			line: node.loc.start.line,
+			column: node.loc.start.column,
+		};
+	}
+	return undefined;
+}
+
+export function createRuntimeCall(fnName: string, args: t.Expression[]): t.AwaitExpression {
+	return t.awaitExpression(
+		t.callExpression(t.memberExpression(t.identifier('__runtime'), t.identifier(fnName)), args)
+	);
+}
+
+export function wrapInAsyncFunction(body: t.Statement[]): t.FunctionExpression {
+	return t.functionExpression(null, [], t.blockStatement(body), false, true);
+}
+
+export function isArrayMethod(node: t.Node, methodName: string): boolean {
+	if (!t.isCallExpression(node)) {
+		return false;
+	}
+
+	const callee = node.callee;
+	if (!t.isMemberExpression(callee)) {
+		return false;
+	}
+
+	return t.isIdentifier(callee.property) && callee.property.name === methodName;
+}

package/src/types.ts

@@ -0,0 +1,101 @@
+import type { Node } from '@babel/types';
+import type { CacheProvider } from '@mondaydotcomorg/atp-protocol';
+
+export interface DetectionResult {
+	needsTransform: boolean;
+	patterns: AsyncPattern[];
+	batchableParallel?: boolean;
+}
+
+export type AsyncPattern =
+	| 'for-of-await'
+	| 'while-await'
+	| 'map-async'
+	| 'forEach-async'
+	| 'filter-async'
+	| 'reduce-async'
+	| 'find-async'
+	| 'some-async'
+	| 'every-async'
+	| 'flatMap-async'
+	| 'promise-all'
+	| 'promise-allSettled';
+
+export interface TransformResult {
+	code: string;
+	transformed: boolean;
+	patterns: AsyncPattern[];
+	metadata: TransformMetadata;
+}
+
+export interface TransformMetadata {
+	loopCount: number;
+	arrayMethodCount: number;
+	parallelCallCount: number;
+	batchableCount: number;
+}
+
+export interface LoopCheckpoint {
+	loopId: string;
+	currentIndex: number;
+	results?: unknown[];
+	accumulator?: unknown;
+	completed?: Set<number>;
+	timestamp: number;
+}
+
+export interface BatchCallInfo {
+	type: 'llm' | 'approval' | 'embedding';
+	operation: string;
+	payload: Record<string, unknown>;
+}
+
+export interface RuntimeContext {
+	executionId: string;
+	cache?: CacheProvider;
+	checkpointPrefix?: string;
+}
+
+export interface TransformerOptions {
+	generateDebugInfo?: boolean;
+	maxLoopNesting?: number;
+	enableBatchParallel?: boolean;
+	batchSizeThreshold?: number;
+}
+
+export interface PausableCallPattern {
+	namespace: string;
+	method: string;
+}
+
+export const PAUSABLE_CALL_PATTERNS: PausableCallPattern[] = [
+	{ namespace: 'atp.llm', method: 'call' },
+	{ namespace: 'atp.llm', method: 'extract' },
+	{ namespace: 'atp.llm', method: 'classify' },
+	{ namespace: 'atp.llm', method: 'stream' },
+	{ namespace: 'atp.llm', method: 'generate' },
+	{ namespace: 'atp.approval', method: 'request' },
+	{ namespace: 'atp.approval', method: 'confirm' },
+	{ namespace: 'atp.approval', method: 'verify' },
+	{ namespace: 'atp.embedding', method: 'embed' },
+	{ namespace: 'atp.embedding', method: 'search' },
+	{ namespace: 'atp.embedding', method: 'create' },
+	{ namespace: 'atp.embedding', method: 'generate' },
+	{ namespace: 'atp.embedding', method: 'encode' },
+];
+
+export interface CompilerConfig {
+	enableBatchParallel?: boolean;
+	maxLoopNesting?: number;
+	checkpointInterval?: number;
+	debugMode?: boolean;
+	batchSizeThreshold?: number;
+}
+
+export const DEFAULT_COMPILER_CONFIG: CompilerConfig = {
+	enableBatchParallel: true,
+	maxLoopNesting: 10,
+	checkpointInterval: 1,
+	debugMode: false,
+	batchSizeThreshold: 10,
+};

package/tsconfig.json

@@ -0,0 +1,12 @@
+{
+	"extends": "../../tsconfig.json",
+	"compilerOptions": {
+		"outDir": "./dist",
+		"declaration": true,
+		"declarationMap": true,
+		"sourceMap": true,
+		"skipLibCheck": true
+	},
+	"include": ["src/**/*"],
+	"exclude": ["node_modules", "dist", "__tests__"]
+}