Setting the file. One moment. scripts/index.ts
TypeScript·392 lines·10 KB
9 buildBatchPrompt,
10 unpackBatchResults,
11} from "./batching.js";
12import type {
13 CreateSource,
14 SwarmHandle,
15 RunOptions,
16 RunResult,
17 RowsOptions,
18 TaskSpec,
19 TaskResult,
20} from "./types.js";
21
22/**
23 * Maximum concurrent subagent dispatches per `run()` call.
24 *
25 * When matched rows exceed this and no explicit `batchSize` is set,
26 * auto-batching groups rows to stay within this concurrency budget.
27 */
28const MAX_SUBAGENTS = 10;
29
30/**
31 * A dispatch unit is a single task for the executor. It tracks
32 * whether it covers one row (single) or multiple (batch) so the
33 * merge step knows how to unpack the result.
34 */
35interface DispatchUnit {
36 /**
37 * The task to dispatch to the executor.
38 */
39 task: TaskSpec;
40
41 /**
42 * Row IDs covered by this task. Single: length 1. Batch: length > 1.
43 */
44 rowIds: string[];
45}
46
47/**
48 * Build dispatch units from pre-grouped batches.
49 *
50 * Single-row batches produce interpolated per-row prompts with the
51 * user's responseSchema. Multi-row batches produce batch prompts
52 * with a wrapped schema.
53 */
54function buildDispatchUnits(
55 batches: Record<string, unknown>[][],
56 opts: {
57 instruction: string;
58 context?: string;
59 subagentType?: string;
60 responseSchema: Record<string, unknown>;
61 mode: "agent" | "invoke";
62 },
63): { units: DispatchUnit[]; errors: TaskResult[] } {
64 const units: DispatchUnit[] = [];
65 const errors: TaskResult[] = [];
66
67 let batchIndex = 0;
68 for (const batch of batches) {
69 if (batch.length === 1) {
70 // Single-row dispatch: interpolate instruction, use schema directly
71 const row = batch[0];
72 const rowId = String(row.id);
73
74 try {
75 let prompt = interpolate(opts.instruction, row);
76 if (opts.context) {
77 prompt = `${opts.context}\n\n${prompt}`;
78 }
79
80 units.push({
81 task: {
82 id: rowId,
83 prompt,
84 subagentType: opts.subagentType,
85 responseSchema: opts.responseSchema,
86 mode: opts.mode,
87 },
88 rowIds: [rowId],
89 });
90 } catch (err) {
91 errors.push({
92 id: rowId,
93 status: "failed",
94 error: (err as Error).message,
95 });
96 }
97 } else {
98 // Multi-row batch: build batch prompt, wrap schema
99 const rowIds = batch.map((r) => String(r.id));
100 units.push({
101 task: {
102 id: `batch_${batchIndex}`,
103 prompt: buildBatchPrompt(opts.instruction, batch, opts.context),
104 subagentType: opts.subagentType,
105 responseSchema: wrapSchema(opts.responseSchema, batch.length),
106 mode: opts.mode,
107 },
108 rowIds,
109 });
110 batchIndex++;
111 }
112 }
113
114 return { units, errors };
115}
116
117/**
118 * Normalize dispatch results into per-row results.
119 *
120 * Single-row units pass through directly. Batch units are unpacked
121 * into one result per row — missing rows become failures.
122 */
123function unpackDispatchResults(
124 units: DispatchUnit[],
125 results: TaskResult[],
126): TaskResult[] {
127 const rowResults: TaskResult[] = [];
128
129 for (let idx = 0; idx < units.length; idx++) {
130 const unit = units[idx];
131 const result = results[idx];
132
133 if (unit.rowIds.length === 1) {
134 rowResults.push(result);
135 continue;
136 }
137
138 if (result.status === "failed") {
139 for (const rowId of unit.rowIds) {
140 rowResults.push({ id: rowId, status: "failed", error: result.error });
141 }
142 continue;
143 }
144
145 const { results: unpacked } = unpackBatchResults(
146 result.result ?? "",
147 unit.rowIds,
148 );
149 for (const rowId of unit.rowIds) {
150 const value = unpacked.get(rowId);
151 if (value !== undefined) {
152 rowResults.push({
153 id: rowId,
154 status: "completed",
155 result: typeof value === "string" ? value : JSON.stringify(value),
156 });
157 } else {
158 rowResults.push({
159 id: rowId,
160 status: "failed",
161 error: "Missing from batch response",
162 });
163 }
164 }
165 }
166
167 return rowResults;
168}
169
170/**
171 * Parse and merge per-row results into table rows.
172 *
173 * Each completed result is JSON-parsed and spread onto the
174 * corresponding row via `mergeResult`.
175 */
176function mergeRowResults(
177 rowResults: TaskResult[],
178 rowById: Map<string, Record<string, unknown>>,
179): { completed: number; failed: number } {
180 let completed = 0;
181 let failed = 0;
182
183 for (const result of rowResults) {
184 const row = rowById.get(result.id);
185 if (!row) {
186 failed++;
187 continue;
188 }
189
190 if (result.status === "completed" && result.result != null) {
191 try {
192 mergeResult(row, JSON.parse(result.result));
193 completed++;
194 } catch {
195 failed++;
196 }
197 } else {
198 failed++;
199 }
200 }
201
202 return { completed, failed };
203}
204
205/**
206 * Verify every `{column}` reference in `instruction` resolves on at
207 * least one matched row. Throws with a list of unresolved paths.
208 */
209function validatePlaceholders(
210 instruction: string,
211 rows: Record<string, unknown>[],
212): void {
213 const placeholders = extractPlaceholders(instruction);
214 if (placeholders.length === 0) {
215 return;
216 }
217 const unresolved = placeholders.filter(
218 (p) => !rows.some((r) => readColumn(r, p) !== undefined),
219 );
220 if (unresolved.length > 0) {
221 throw new Error(
222 `instruction references unknown column(s): ${unresolved.join(", ")}`,
223 );
224 }
225}
226
227/**
228 * Create a table from a source specification and persist it to the backend.
229 *
230 * Thin wrapper around `createTable` — validates the source, builds rows,
231 * runs eviction if necessary, and persists the table as JSONL.
232 *
233 * @param source - Exactly one of `glob`, `filePaths`, or `tasks`.
234 * @returns A lightweight handle with the table's ID, row count, and columns.
235 */
236export async function create(source: CreateSource): Promise<SwarmHandle> {
237 return createTable(source);
238}
239
240/**
241 * Dispatch work across table rows and update the table in place.
242 *
243 * Loads the table, partitions rows by filter, interpolates the
244 * instruction template per-row (or builds batch prompts), dispatches
245 * to subagents via `tools.swarm_task()`, merges results into rows,
246 * and persists the updated table.
247 *
248 * @param handle - A table handle or object with an `id` field.
249 * @param options - Dispatch configuration (instruction, filter, schema, etc.).
250 * @returns A summary with completion counts and deduplicated failure groups.
251 */
252export async function run(
253 tableId: string,
254 options: RunOptions,
255): Promise<RunResult> {
256 const allRows = await loadTable(tableId);
257 const {
258 instruction,
259 context,
260 filter,
261 subagentType,
262 responseSchema,
263 batchSize,
264 concurrency,
265 } = options;
266 const mode = subagentType != null ? "agent" : "invoke";
267
268 const effectiveConcurrency = Math.max(
269 1,
270 Math.min(concurrency ?? MAX_SUBAGENTS, MAX_SUBAGENTS),
271 );
272
273 // -----------------------------------------------------------------------
274 // 1. Partition rows into matched (dispatched) and skipped (filtered out)
275 // -----------------------------------------------------------------------
276
277 const matched: Record<string, unknown>[] = [];
278 let skippedCount = 0;
279
280 for (const row of allRows) {
281 if (!filter || evaluateFilter(filter, row)) {
282 matched.push(row);
283 } else {
284 skippedCount++;
285 }
286 }
287
288 if (matched.length === 0) {
289 return {
290 completed: 0,
291 failed: 0,
292 skipped: allRows.length,
293 failures: [],
294 };
295 }
296
297 validatePlaceholders(instruction, matched);
298
299 // -----------------------------------------------------------------------
300 // 2. Resolve batches and build dispatch units
301 // -----------------------------------------------------------------------
302
303 const batches = resolveBatchGroups(matched, effectiveConcurrency, batchSize);
304
305 const { units, errors: interpolationErrors } = buildDispatchUnits(batches, {
306 instruction,
307 context,
308 subagentType,
309 responseSchema,
310 mode,
311 });
312
313 // -----------------------------------------------------------------------
314 // 3. Dispatch
315 // -----------------------------------------------------------------------
316
317 const dispatchResults = await dispatch(
318 units.map((u) => u.task),
319 { concurrency: effectiveConcurrency },
320 );
321
322 // -----------------------------------------------------------------------
323 // 4. Unpack and merge results into rows
324 // -----------------------------------------------------------------------
325
326 const rowById = new Map<string, Record<string, unknown>>();
327 for (const row of matched) {
328 rowById.set(String(row.id), row);
329 }
330
331 const rowResults = unpackDispatchResults(units, dispatchResults);
332 const { completed, failed: mergeFailed } = mergeRowResults(
333 rowResults,
334 rowById,
335 );
336 const failed = mergeFailed + interpolationErrors.length;
337 const allRowResults = [...interpolationErrors, ...rowResults];
338
339 // -----------------------------------------------------------------------
340 // 5. Persist and return summary
341 // -----------------------------------------------------------------------
342
343 await saveTable(tableId, allRows);
344
345 return {
346 completed,
347 failed,
348 skipped: skippedCount,
349 failures: deduplicateFailures(allRowResults),
350 };
351}
352
353/**
354 * Retrieve rows from a table, optionally filtered and projected.
355 *
356 * Loads the table and applies filter, column projection, and row
357 * limiting in that order. Use for inspection and JS-based aggregation
358 * — the heavy data stays in the sandbox and only the computed result
359 * (via `console.log`) goes back to the agent's context.
360 *
361 * @param handle - A table handle or object with an `id` field.
362 * @param options - Optional filtering, projection, and limiting.
363 * @returns Array of row objects matching the criteria.
364 */
365export async function rows(
366 tableId: string,
367 options?: RowsOptions,
368): Promise<Record<string, unknown>[]> {
369 let result = await loadTable(tableId);
370
371 if (options?.filter) {
372 const f = options.filter;
373 result = result.filter((row) => evaluateFilter(f, row));
374 }
375
376 if (options?.columns) {
377 const cols = options.columns;
378 result = result.map((row) => {
379 const projected: Record<string, unknown> = {};
380 for (const col of cols) {
381 if (col in row) projected[col] = row[col];
382 }
383 return projected;
384 });
385 }
386
387 if (options?.limit != null && options.limit >= 0) {
388 result = result.slice(0, options.limit);
389 }
390
391 return result;
392}