@mondaydotcomorg/atp-compiler 0.19.23 → 0.19.26

package/src/transformer/array-transformer-batch-reconstruct.ts

@@ -3,14 +3,15 @@ import _traverse from '@babel/traverse';
 const traverse = typeof (_traverse as any).default === 'function' ? (_traverse as any).default : _traverse;
 import { generateUniqueId } from '../runtime/context.js';
 import { BatchParallelDetector } from './batch-detector.js';
-import { findLLMCallExpression } from './array-transformer-utils.js';
+import { findAllLLMCallExpressions } from './array-transformer-utils.js';
 
 /**
  * Transform array method to batch LLM calls while preserving callback logic.
- * 
- * This creates a two-step transformation:
- * 1. Batch all LLM calls in parallel
- * 2. Reconstruct objects using the batched results 
+ * Supports multiple LLM calls per callback.
+ *
+ * This creates a multi-step transformation:
+ * 1. Batch all LLM calls in parallel (one batch per unique type:operation)
+ * 2. Reconstruct objects using the batched results
  */
 export function transformToBatchWithReconstruction(
 	path: any,
@@ -28,74 +29,117 @@ export function transformToBatchWithReconstruction(
 	if (!t.isIdentifier(paramName)) {
 		return false;
 	}
-
 	const param = paramName.name;
 	const array = (node.callee as t.MemberExpression).object;
 
-	const llmCall = findLLMCallExpression(callback.body);
-	if (!llmCall) {
+	// Find ALL LLM calls
+	const llmCalls = findAllLLMCallExpressions(callback.body, batchDetector);
+	if (llmCalls.length === 0) {
 		return false;
 	}
 
-	const callInfo = batchDetector.extractCallInfo(llmCall);
-	if (!callInfo) {
-		return false;
-	}
+	const methodId = generateUniqueId(`${methodName}_batch_reconstruct`);
 
-	const payloadNode = batchDetector.extractPayloadNode(llmCall);
-	if (!payloadNode) {
-		return false;
-	}
+	const originalIndexParam = callback.params[1];
+	const indexVar =
+		originalIndexParam && t.isIdentifier(originalIndexParam) ? originalIndexParam.name : '__idx';
+
+	// Create batch declarations - one per LLM call (in order of appearance)
+	// This ensures each call gets its own result array
+	const batchDeclarations: t.Statement[] = [];
+	const resultVarByCallIndex = new Map<number, string>();
+
+	for (let i = 0; i < llmCalls.length; i++) {
+		const call = llmCalls[i]!;
+		const resultsVar = `__batch_results_${i}_${methodId.replace(/[^a-zA-Z0-9]/g, '_')}`;
+		resultVarByCallIndex.set(i, resultsVar);
+
+		const payloadMapper = t.arrowFunctionExpression(
+			[t.identifier(param)],
+			t.objectExpression([
+				t.objectProperty(t.identifier('type'), t.stringLiteral(call.callInfo.type)),
+				t.objectProperty(
+					t.identifier('operation'),
+					t.stringLiteral(call.callInfo.operation)
+				),
+				t.objectProperty(t.identifier('payload'), t.cloneNode(call.payloadNode, true)),
+			])
+		);
 
-	const methodId = generateUniqueId(`${methodName}_batch_reconstruct`);
-	const resultsVar = `__batch_results_${methodId.replace(/[^a-zA-Z0-9]/g, '_')}`;
-	const indexVar = '__idx';
-
-	const payloadMapper = t.arrowFunctionExpression(
-		[t.identifier(param)],
-		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'), t.cloneNode(payloadNode, true)),
-		])
-	);
+		const batchCall = t.awaitExpression(
+			t.callExpression(
+				t.memberExpression(t.identifier('__runtime'), t.identifier('batchParallel')),
+				[
+					t.callExpression(
+						t.memberExpression(t.cloneNode(array, true), t.identifier('map')),
+						[payloadMapper]
+					),
+					t.stringLiteral(`${methodId}_${i}`),
+				]
+			)
+		);
 
+		batchDeclarations.push(
+			t.variableDeclaration('const', [
+				t.variableDeclarator(t.identifier(resultsVar), batchCall),
+			])
+		);
+	}
+
+	// Clone the callback body for reconstruction
 	const clonedBody = t.cloneNode(callback.body, true);
-	
-	const resultAccess = t.memberExpression(
-		t.identifier(resultsVar),
-		t.identifier(indexVar),
-		true
-	);
 
 	let traversableNode: t.Statement;
 	if (t.isBlockStatement(clonedBody)) {
-		traversableNode = t.functionDeclaration(
-			t.identifier('__temp'),
-			[],
-			clonedBody
-		);
+		traversableNode = t.functionDeclaration(t.identifier('__temp'), [], clonedBody);
 	} else {
 		traversableNode = t.expressionStatement(clonedBody as t.Expression);
 	}
 
-	let replaced = false;
+	// Replace each await expression with the corresponding result access
+	// We match calls by comparing their structure to the original calls
+	let replacementCount = 0;
+
 	traverse(t.file(t.program([traversableNode])), {
 		AwaitExpression(awaitPath: any) {
-			if (replaced) return;
 			const arg = awaitPath.node.argument;
-			if (t.isCallExpression(arg)) {
-				const info = batchDetector.extractCallInfo(arg);
-				if (info && info.type === callInfo.type && info.operation === callInfo.operation) {
-					awaitPath.replaceWith(resultAccess);
-					replaced = true;
+			if (!t.isCallExpression(arg)) return;
+
+			const info = batchDetector.extractCallInfo(arg);
+			if (!info) return;
+
+			const key = `${info.type}:${info.operation}`;
+
+			// Find first unused call with matching key
+			let matchedIndex = -1;
+			for (let i = 0; i < llmCalls.length; i++) {
+				const original = llmCalls[i]!;
+				if (original.key === key && resultVarByCallIndex.has(i)) {
+					matchedIndex = i;
+					break;
 				}
 			}
+
+			if (matchedIndex === -1) return;
+
+			const resultsVar = resultVarByCallIndex.get(matchedIndex);
+			if (!resultsVar) return;
+			// Remove from map so we don't reuse it
+			resultVarByCallIndex.delete(matchedIndex);
+
+			const resultAccess = t.memberExpression(
+				t.identifier(resultsVar),
+				t.identifier(indexVar),
+				true
+			);
+
+			awaitPath.replaceWith(resultAccess);
+			replacementCount++;
 		},
 		noScope: true,
 	});
 
-	if (!replaced) {
+	if (replacementCount === 0) {
 		return false;
 	}
 
@@ -112,36 +156,17 @@ export function transformToBatchWithReconstruction(
 	);
 	reconstructMapper.async = false;
 
-	const batchCall = t.awaitExpression(
-		t.callExpression(
-			t.memberExpression(t.identifier('__runtime'), t.identifier('batchParallel')),
-			[
-				t.callExpression(
-					t.memberExpression(t.cloneNode(array, true), t.identifier('map')),
-					[payloadMapper]
-				),
-				t.stringLiteral(methodId),
-			]
-		)
-	);
-
-	const resultsDeclaration = t.variableDeclaration('const', [
-		t.variableDeclarator(t.identifier(resultsVar), batchCall),
-	]);
-
 	const reconstructCall = t.callExpression(
 		t.memberExpression(t.cloneNode(array, true), t.identifier('map')),
 		[reconstructMapper]
 	);
 
+	// Build the IIFE with all batch declarations followed by reconstruction
 	const iife = t.callExpression(
 		t.arrowFunctionExpression(
 			[],
-			t.blockStatement([
-				resultsDeclaration,
-				t.returnStatement(reconstructCall),
-			]),
-			true
+			t.blockStatement([...batchDeclarations, t.returnStatement(reconstructCall)]),
+			true // async
 		),
 		[]
 	);
@@ -152,4 +177,3 @@ export function transformToBatchWithReconstruction(
 	onTransform();
 	return true;
 }
-

package/src/transformer/array-transformer-utils.ts

@@ -1,23 +1,168 @@
 import * as t from '@babel/types';
 import { isArrayMethod } from './utils.js';
+import type { BatchParallelDetector } from './batch-detector.js';
+import type { BatchCallInfo } from '../types.js';
+
+export interface LLMCallInfo {
+	callNode: t.CallExpression;
+	callInfo: BatchCallInfo;
+	payloadNode: t.Expression;
+	// Unique key for grouping: "type:operation"
+	key: string;
+}
 
 /**
- * Find LLM call expression in AST node
+ * Collect all identifiers referenced in an expression
  */
-export function findLLMCallExpression(body: t.Node): t.CallExpression | null {
-	let found: t.CallExpression | null = null;
+function collectReferencedIdentifiers(node: t.Node): Set<string> {
+	const identifiers = new Set<string>();
+
+	const visit = (n: t.Node) => {
+		if (t.isIdentifier(n)) {
+			identifiers.add(n.name);
+		}
+
+		// Continue traversing
+		Object.keys(n).forEach((key) => {
+			// Skip 'type' and other metadata fields
+			if (key === 'type' || key === 'loc' || key === 'start' || key === 'end') return;
+			const value = (n as any)[key];
+			if (Array.isArray(value)) {
+				value.forEach((item) => {
+					if (item && typeof item === 'object' && item.type) {
+						visit(item);
+					}
+				});
+			} else if (value && typeof value === 'object' && value.type) {
+				visit(value);
+			}
+		});
+	};
+
+	visit(node);
+	return identifiers;
+}
+
+/**
+ * Collect all variable names declared inside a callback body.
+ * This includes: const/let/var declarations, function parameters, etc.
+ */
+function collectLocalVariables(body: t.Node): Set<string> {
+	const locals = new Set<string>();
 
 	const visit = (node: t.Node) => {
-		if (found) return;
+		// Variable declarations: const x = ..., let y = ..., var z = ...
+		if (t.isVariableDeclaration(node)) {
+			for (const decl of node.declarations) {
+				if (t.isIdentifier(decl.id)) {
+					locals.add(decl.id.name);
+				} else if (t.isObjectPattern(decl.id)) {
+					// Destructuring: const { a, b } = ...
+					for (const prop of decl.id.properties) {
+						if (t.isObjectProperty(prop) && t.isIdentifier(prop.value)) {
+							locals.add(prop.value.name);
+						} else if (t.isRestElement(prop) && t.isIdentifier(prop.argument)) {
+							locals.add(prop.argument.name);
+						}
+					}
+				} else if (t.isArrayPattern(decl.id)) {
+					// Destructuring: const [a, b] = ...
+					for (const elem of decl.id.elements) {
+						if (t.isIdentifier(elem)) {
+							locals.add(elem.name);
+						}
+					}
+				}
+			}
+		}
 
+		// Continue traversing (but don't descend into nested functions - their locals are their own scope)
+		if (t.isFunction(node)) {
+			return; // Don't traverse into nested functions
+		}
+
+		Object.keys(node).forEach((key) => {
+			if (key === 'type' || key === 'loc' || key === 'start' || key === 'end') return;
+			const value = (node as any)[key];
+			if (Array.isArray(value)) {
+				value.forEach((item) => {
+					if (item && typeof item === 'object' && item.type) {
+						visit(item);
+					}
+				});
+			} else if (value && typeof value === 'object' && value.type) {
+				visit(value);
+			}
+		});
+	};
+
+	visit(body);
+	return locals;
+}
+
+/**
+ * Check if LLM call payloads depend on variables that are local to the callback.
+ * This includes:
+ * - Variables computed from other expressions (const x = item.name.toUpperCase())
+ * - Results from previous LLM calls (const title = await atp.llm.call(...))
+ * - Any other locally-defined variable
+ *
+ * When batch transforming, the payload mapper only has access to the array item parameter
+ * and outer scope variables - NOT local variables defined inside the callback.
+ *
+ * @param body - The callback body
+ * @param itemParamName - The name of the array item parameter (e.g., "item")
+ * @param batchDetector - The batch detector for extracting LLM call info
+ * @returns true if there are dependencies that prevent batch transformation
+ */
+export function hasLLMCallDependencies(
+	body: t.Node,
+	batchDetector: BatchParallelDetector,
+	itemParamName?: string
+): boolean {
+	// Find all locally-defined variables in the callback body
+	const localVariables = collectLocalVariables(body);
+
+	// If no local variables, no dependencies possible
+	if (localVariables.size === 0) {
+		return false;
+	}
+
+	// Now check if any LLM call payload references these local variables
+	const allCalls = findAllAwaitedMemberCalls(body);
+
+	for (const call of allCalls) {
+		const payloadNode = batchDetector.extractPayloadNode(call);
+		if (payloadNode) {
+			const referencedIds = collectReferencedIdentifiers(payloadNode);
+			// Check if any referenced identifier is a local variable
+			// (excluding the item parameter which is passed to the payload mapper)
+			for (const id of referencedIds) {
+				if (localVariables.has(id) && id !== itemParamName) {
+					return true; // Found a dependency on a local variable
+				}
+			}
+		}
+	}
+
+	return false;
+}
+
+/**
+ * Find all awaited member expression calls in AST node
+ */
+function findAllAwaitedMemberCalls(body: t.Node): t.CallExpression[] {
+	const calls: t.CallExpression[] = [];
+
+	const visit = (node: t.Node) => {
 		if (t.isAwaitExpression(node) && t.isCallExpression(node.argument)) {
 			const call = node.argument;
 			if (t.isMemberExpression(call.callee)) {
-				found = call;
-				return;
+				calls.push(call);
 			}
 		}
 
+		// Continue traversing
 		Object.keys(node).forEach((key) => {
 			const value = (node as any)[key];
 			if (Array.isArray(value)) {
@@ -33,7 +178,42 @@ export function findLLMCallExpression(body: t.Node): t.CallExpression | null {
 	};
 
 	visit(body);
-	return found;
+	return calls;
+}
+
+/**
+ * Find ALL LLM call expressions in AST node with batch info
+ */
+export function findAllLLMCallExpressions(
+	body: t.Node,
+	batchDetector: BatchParallelDetector
+): LLMCallInfo[] {
+	const allCalls = findAllAwaitedMemberCalls(body);
+	const llmCalls: LLMCallInfo[] = [];
+
+	for (const call of allCalls) {
+		const callInfo = batchDetector.extractCallInfo(call);
+		const payloadNode = batchDetector.extractPayloadNode(call);
+
+		if (callInfo && payloadNode) {
+			llmCalls.push({
+				callNode: call,
+				callInfo,
+				payloadNode,
+				key: `${callInfo.type}:${callInfo.operation}`,
+			});
+		}
+	}
+
+	return llmCalls;
+}
+
+/**
+ * Find first LLM call expression in AST node
+ */
+export function findLLMCallExpression(body: t.Node): t.CallExpression | null {
+	const calls = findAllAwaitedMemberCalls(body);
+	return calls[0] ?? null;
 }
 
 /**

package/src/transformer/array-transformer.ts

@@ -5,6 +5,7 @@ import {
 	getArrayMethodName,
 	canUseBatchParallel,
 	findLLMCallExpression,
+	hasLLMCallDependencies,
 } from './array-transformer-utils.js';
 import { transformToBatchParallel } from './array-transformer-batch.js';
 import { transformToSequential } from './array-transformer-sequential.js';
@@ -36,15 +37,29 @@ export class ArrayTransformer {
 		}
 
 		const batchResult = this.batchOptimizer.canBatchArrayMethod(callback);
-
 		if (!batchResult.canBatch && methodName === 'map') {
 			const reason = batchResult.reason || '';
-			const hasObjectOrArrayReturn =
-				reason.includes('object expression') || reason.includes('array expression');
+			// Try batch-with-reconstruction for:
+			// 1. Object/array returns (would lose structure with simple batch)
+			// 2. Multiple pausable calls (can batch each call type separately)
+			const canTryBatchReconstruct =
+				reason.includes('object expression') ||
+				reason.includes('array expression') ||
+				reason.includes('Multiple pausable calls');
 
-			if (hasObjectOrArrayReturn) {
+			if (canTryBatchReconstruct) {
 				const llmCall = findLLMCallExpression(callback.body);
-				if (llmCall) {
+				// Get the item parameter name (e.g., "item" in items.map(async (item) => ...))
+				const itemParam = callback.params[0];
+				const itemParamName = t.isIdentifier(itemParam) ? itemParam.name : undefined;
+				// Check for dependencies on local variables (e.g., computed values, previous LLM results)
+				// If dependencies exist, we can't batch because payload mapper only has access to item + outer scope
+				const hasDependencies = hasLLMCallDependencies(
+					callback.body,
+					this.batchDetector,
+					itemParamName
+				);
+				if (llmCall && !hasDependencies) {
 					const success = transformToBatchWithReconstruction(
 						path,
 						node,