ak-gemini 2.0.7 → 2.0.9

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/GUIDE.md CHANGED
@@ -349,6 +349,22 @@ console.log(result.text); // "There were 47 new signups this week. I've se
349
349
  console.log(result.toolCalls); // [{ name: 'query_db', args: {...}, result: [...] }, { name: 'send_email', ... }]
350
350
  ```
351
351
 
352
+ ### Parallel Tool Execution
353
+
354
+ Control whether tool calls within a round execute in parallel or sequentially:
355
+
356
+ ```javascript
357
+ const agent = new ToolAgent({
358
+ tools: [...],
359
+ toolExecutor: myExecutor,
360
+ parallelToolCalls: true, // default: unlimited parallel execution
361
+ // parallelToolCalls: false, // sequential execution
362
+ // parallelToolCalls: 3, // max 3 concurrent tool executions
363
+ });
364
+ ```
365
+
366
+ When the model returns multiple tool calls in a single response, parallel execution runs them concurrently — significantly faster for I/O-bound tools (HTTP requests, database queries, etc.).
367
+
352
368
  ### Streaming
353
369
 
354
370
  Stream the agent's output in real-time — useful for showing progress in a UI:
package/README.md CHANGED
@@ -413,6 +413,7 @@ All classes accept `BaseGeminiOptions`:
413
413
  | `maxToolRounds` | number | `10` | Max tool-use loop iterations |
414
414
  | `onToolCall` | function | — | Notification callback when tool is called |
415
415
  | `onBeforeExecution` | function | — | `async (toolName, args) => boolean` — gate execution |
416
+ | `parallelToolCalls` | boolean \| number | `true` | Parallel tool execution: `false` = sequential, `true` = unlimited, number = concurrency limit |
416
417
 
417
418
  ### CodeAgent-Specific
418
419
 
package/index.cjs CHANGED
@@ -1325,6 +1325,24 @@ var Message = class extends base_default {
1325
1325
  var message_default = Message;
1326
1326
 
1327
1327
  // tool-agent.js
1328
+ async function runWithConcurrency(tasks, concurrency) {
1329
+ if (concurrency === Infinity) return Promise.all(tasks.map((t) => t()));
1330
+ if (concurrency === 1) {
1331
+ const results2 = [];
1332
+ for (const t of tasks) results2.push(await t());
1333
+ return results2;
1334
+ }
1335
+ const results = new Array(tasks.length);
1336
+ let next = 0;
1337
+ async function worker() {
1338
+ while (next < tasks.length) {
1339
+ const i = next++;
1340
+ results[i] = await tasks[i]();
1341
+ }
1342
+ }
1343
+ await Promise.all(Array.from({ length: Math.min(concurrency, tasks.length) }, () => worker()));
1344
+ return results;
1345
+ }
1328
1346
  var ToolAgent = class extends base_default {
1329
1347
  /**
1330
1348
  * @param {ToolAgentOptions} [options={}]
@@ -1342,6 +1360,8 @@ var ToolAgent = class extends base_default {
1342
1360
  if (this.toolExecutor && this.tools.length === 0) {
1343
1361
  throw new Error("ToolAgent: toolExecutor provided without tools. Provide tool declarations so the model knows what tools are available.");
1344
1362
  }
1363
+ this.parallelToolCalls = options.parallelToolCalls ?? true;
1364
+ this._concurrency = this.parallelToolCalls === true ? Infinity : this.parallelToolCalls === false ? 1 : this.parallelToolCalls;
1345
1365
  this.maxToolRounds = options.maxToolRounds || 10;
1346
1366
  this.onToolCall = options.onToolCall || null;
1347
1367
  this.onBeforeExecution = options.onBeforeExecution || null;
@@ -1372,38 +1392,36 @@ var ToolAgent = class extends base_default {
1372
1392
  if (this._stopped) break;
1373
1393
  const functionCalls = response.functionCalls;
1374
1394
  if (!functionCalls || functionCalls.length === 0) break;
1375
- const toolResults = await Promise.all(
1376
- functionCalls.map(async (call) => {
1377
- if (this.onToolCall) {
1378
- try {
1379
- this.onToolCall(call.name, call.args);
1380
- } catch (e) {
1381
- logger_default.warn(`onToolCall callback error: ${e.message}`);
1382
- }
1383
- }
1384
- if (this.onBeforeExecution) {
1385
- try {
1386
- const allowed = await this.onBeforeExecution(call.name, call.args);
1387
- if (allowed === false) {
1388
- const result2 = { error: "Execution denied by onBeforeExecution callback" };
1389
- allToolCalls.push({ name: call.name, args: call.args, result: result2 });
1390
- return { id: call.id, name: call.name, result: result2 };
1391
- }
1392
- } catch (e) {
1393
- logger_default.warn(`onBeforeExecution callback error: ${e.message}`);
1394
- }
1395
+ const tasks = functionCalls.map((call) => async () => {
1396
+ if (this.onToolCall) {
1397
+ try {
1398
+ this.onToolCall(call.name, call.args);
1399
+ } catch (e) {
1400
+ logger_default.warn(`onToolCall callback error: ${e.message}`);
1395
1401
  }
1396
- let result;
1402
+ }
1403
+ if (this.onBeforeExecution) {
1397
1404
  try {
1398
- result = await this.toolExecutor(call.name, call.args);
1399
- } catch (err) {
1400
- logger_default.warn(`Tool ${call.name} failed: ${err.message}`);
1401
- result = { error: err.message };
1405
+ const allowed = await this.onBeforeExecution(call.name, call.args);
1406
+ if (allowed === false) {
1407
+ const result2 = { error: "Execution denied by onBeforeExecution callback" };
1408
+ return { id: call.id, name: call.name, args: call.args, result: result2 };
1409
+ }
1410
+ } catch (e) {
1411
+ logger_default.warn(`onBeforeExecution callback error: ${e.message}`);
1402
1412
  }
1403
- allToolCalls.push({ name: call.name, args: call.args, result });
1404
- return { id: call.id, name: call.name, result };
1405
- })
1406
- );
1413
+ }
1414
+ let result;
1415
+ try {
1416
+ result = await this.toolExecutor(call.name, call.args);
1417
+ } catch (err) {
1418
+ logger_default.warn(`Tool ${call.name} failed: ${err.message}`);
1419
+ result = { error: err.message };
1420
+ }
1421
+ return { id: call.id, name: call.name, args: call.args, result };
1422
+ });
1423
+ const toolResults = await runWithConcurrency(tasks, this._concurrency);
1424
+ for (const r of toolResults) allToolCalls.push({ name: r.name, args: r.args, result: r.result });
1407
1425
  response = await this._withRetry(() => this.chatSession.sendMessage({
1408
1426
  message: toolResults.map((r) => ({
1409
1427
  functionResponse: {
@@ -1471,39 +1489,81 @@ var ToolAgent = class extends base_default {
1471
1489
  return;
1472
1490
  }
1473
1491
  const toolResults = [];
1474
- for (const call of functionCalls) {
1475
- if (this._stopped) break;
1476
- yield { type: "tool_call", toolName: call.name, args: call.args };
1477
- if (this.onToolCall) {
1478
- try {
1479
- this.onToolCall(call.name, call.args);
1480
- } catch (e) {
1481
- logger_default.warn(`onToolCall callback error: ${e.message}`);
1492
+ if (this._concurrency === 1) {
1493
+ for (const call of functionCalls) {
1494
+ if (this._stopped) break;
1495
+ yield { type: "tool_call", toolName: call.name, args: call.args };
1496
+ if (this.onToolCall) {
1497
+ try {
1498
+ this.onToolCall(call.name, call.args);
1499
+ } catch (e) {
1500
+ logger_default.warn(`onToolCall callback error: ${e.message}`);
1501
+ }
1482
1502
  }
1483
- }
1484
- let denied = false;
1485
- if (this.onBeforeExecution) {
1486
- try {
1487
- const allowed = await this.onBeforeExecution(call.name, call.args);
1488
- if (allowed === false) denied = true;
1489
- } catch (e) {
1490
- logger_default.warn(`onBeforeExecution callback error: ${e.message}`);
1503
+ let denied = false;
1504
+ if (this.onBeforeExecution) {
1505
+ try {
1506
+ const allowed = await this.onBeforeExecution(call.name, call.args);
1507
+ if (allowed === false) denied = true;
1508
+ } catch (e) {
1509
+ logger_default.warn(`onBeforeExecution callback error: ${e.message}`);
1510
+ }
1491
1511
  }
1512
+ let result;
1513
+ if (denied) {
1514
+ result = { error: "Execution denied by onBeforeExecution callback" };
1515
+ } else {
1516
+ try {
1517
+ result = await this.toolExecutor(call.name, call.args);
1518
+ } catch (err) {
1519
+ logger_default.warn(`Tool ${call.name} failed: ${err.message}`);
1520
+ result = { error: err.message };
1521
+ }
1522
+ }
1523
+ allToolCalls.push({ name: call.name, args: call.args, result });
1524
+ yield { type: "tool_result", toolName: call.name, result };
1525
+ toolResults.push({ id: call.id, name: call.name, result });
1492
1526
  }
1493
- let result;
1494
- if (denied) {
1495
- result = { error: "Execution denied by onBeforeExecution callback" };
1496
- } else {
1497
- try {
1498
- result = await this.toolExecutor(call.name, call.args);
1499
- } catch (err) {
1500
- logger_default.warn(`Tool ${call.name} failed: ${err.message}`);
1501
- result = { error: err.message };
1527
+ } else {
1528
+ for (const call of functionCalls) {
1529
+ yield { type: "tool_call", toolName: call.name, args: call.args };
1530
+ }
1531
+ const tasks = functionCalls.map((call) => async () => {
1532
+ if (this.onToolCall) {
1533
+ try {
1534
+ this.onToolCall(call.name, call.args);
1535
+ } catch (e) {
1536
+ logger_default.warn(`onToolCall callback error: ${e.message}`);
1537
+ }
1538
+ }
1539
+ let denied = false;
1540
+ if (this.onBeforeExecution) {
1541
+ try {
1542
+ const allowed = await this.onBeforeExecution(call.name, call.args);
1543
+ if (allowed === false) denied = true;
1544
+ } catch (e) {
1545
+ logger_default.warn(`onBeforeExecution callback error: ${e.message}`);
1546
+ }
1502
1547
  }
1548
+ let result;
1549
+ if (denied) {
1550
+ result = { error: "Execution denied by onBeforeExecution callback" };
1551
+ } else {
1552
+ try {
1553
+ result = await this.toolExecutor(call.name, call.args);
1554
+ } catch (err) {
1555
+ logger_default.warn(`Tool ${call.name} failed: ${err.message}`);
1556
+ result = { error: err.message };
1557
+ }
1558
+ }
1559
+ return { id: call.id, name: call.name, args: call.args, result };
1560
+ });
1561
+ const results = await runWithConcurrency(tasks, this._concurrency);
1562
+ for (const r of results) {
1563
+ allToolCalls.push({ name: r.name, args: r.args, result: r.result });
1564
+ yield { type: "tool_result", toolName: r.name, result: r.result };
1565
+ toolResults.push({ id: r.id, name: r.name, result: r.result });
1503
1566
  }
1504
- allToolCalls.push({ name: call.name, args: call.args, result });
1505
- yield { type: "tool_result", toolName: call.name, result };
1506
- toolResults.push({ id: call.id, name: call.name, result });
1507
1567
  }
1508
1568
  streamResponse = await this._withRetry(() => this.chatSession.sendMessageStream({
1509
1569
  message: toolResults.map((r) => ({
package/package.json CHANGED
@@ -2,7 +2,7 @@
2
2
  "name": "ak-gemini",
3
3
  "author": "ak@mixpanel.com",
4
4
  "description": "AK's Generative AI Helper for doing... everything",
5
- "version": "2.0.7",
5
+ "version": "2.0.9",
6
6
  "main": "index.js",
7
7
  "files": [
8
8
  "index.js",
package/tool-agent.js CHANGED
@@ -13,6 +13,31 @@ import log from './logger.js';
13
13
  * @typedef {import('./types').AgentStreamEvent} AgentStreamEvent
14
14
  */
15
15
 
16
+ /**
17
+ * Execute async task factories with a concurrency limit.
18
+ * @param {Array<() => Promise<any>>} tasks
19
+ * @param {number} concurrency - Infinity for unlimited, 1 for sequential
20
+ * @returns {Promise<any[]>} Results in same order as tasks
21
+ */
22
+ async function runWithConcurrency(tasks, concurrency) {
23
+ if (concurrency === Infinity) return Promise.all(tasks.map(t => t()));
24
+ if (concurrency === 1) {
25
+ const results = [];
26
+ for (const t of tasks) results.push(await t());
27
+ return results;
28
+ }
29
+ const results = new Array(tasks.length);
30
+ let next = 0;
31
+ async function worker() {
32
+ while (next < tasks.length) {
33
+ const i = next++;
34
+ results[i] = await tasks[i]();
35
+ }
36
+ }
37
+ await Promise.all(Array.from({ length: Math.min(concurrency, tasks.length) }, () => worker()));
38
+ return results;
39
+ }
40
+
16
41
  /**
17
42
  * AI agent that uses user-provided tools to accomplish tasks.
18
43
  * Automatically manages the tool-use loop: when the model decides to call
@@ -75,6 +100,13 @@ class ToolAgent extends BaseGemini {
75
100
  throw new Error("ToolAgent: toolExecutor provided without tools. Provide tool declarations so the model knows what tools are available.");
76
101
  }
77
102
 
103
+ // ── Parallel execution ──
104
+ this.parallelToolCalls = options.parallelToolCalls ?? true;
105
+ /** @private */
106
+ this._concurrency = this.parallelToolCalls === true ? Infinity
107
+ : this.parallelToolCalls === false ? 1
108
+ : this.parallelToolCalls;
109
+
78
110
  // ── Tool loop config ──
79
111
  this.maxToolRounds = options.maxToolRounds || 10;
80
112
  this.onToolCall = options.onToolCall || null;
@@ -116,41 +148,39 @@ class ToolAgent extends BaseGemini {
116
148
  const functionCalls = response.functionCalls;
117
149
  if (!functionCalls || functionCalls.length === 0) break;
118
150
 
119
- const toolResults = await Promise.all(
120
- functionCalls.map(async (call) => {
121
- // Fire onToolCall callback
122
- if (this.onToolCall) {
123
- try { this.onToolCall(call.name, call.args); }
124
- catch (e) { log.warn(`onToolCall callback error: ${e.message}`); }
125
- }
151
+ const tasks = functionCalls.map(call => async () => {
152
+ // Fire onToolCall callback
153
+ if (this.onToolCall) {
154
+ try { this.onToolCall(call.name, call.args); }
155
+ catch (e) { log.warn(`onToolCall callback error: ${e.message}`); }
156
+ }
126
157
 
127
- // Check onBeforeExecution gate
128
- if (this.onBeforeExecution) {
129
- try {
130
- const allowed = await this.onBeforeExecution(call.name, call.args);
131
- if (allowed === false) {
132
- const result = { error: 'Execution denied by onBeforeExecution callback' };
133
- allToolCalls.push({ name: call.name, args: call.args, result });
134
- return { id: call.id, name: call.name, result };
135
- }
136
- } catch (e) {
137
- log.warn(`onBeforeExecution callback error: ${e.message}`);
158
+ // Check onBeforeExecution gate
159
+ if (this.onBeforeExecution) {
160
+ try {
161
+ const allowed = await this.onBeforeExecution(call.name, call.args);
162
+ if (allowed === false) {
163
+ const result = { error: 'Execution denied by onBeforeExecution callback' };
164
+ return { id: call.id, name: call.name, args: call.args, result };
138
165
  }
166
+ } catch (e) {
167
+ log.warn(`onBeforeExecution callback error: ${e.message}`);
139
168
  }
169
+ }
140
170
 
141
- let result;
142
- try {
143
- result = await this.toolExecutor(call.name, call.args);
144
- } catch (err) {
145
- log.warn(`Tool ${call.name} failed: ${err.message}`);
146
- result = { error: err.message };
147
- }
171
+ let result;
172
+ try {
173
+ result = await this.toolExecutor(call.name, call.args);
174
+ } catch (err) {
175
+ log.warn(`Tool ${call.name} failed: ${err.message}`);
176
+ result = { error: err.message };
177
+ }
148
178
 
149
- allToolCalls.push({ name: call.name, args: call.args, result });
179
+ return { id: call.id, name: call.name, args: call.args, result };
180
+ });
150
181
 
151
- return { id: call.id, name: call.name, result };
152
- })
153
- );
182
+ const toolResults = await runWithConcurrency(tasks, this._concurrency);
183
+ for (const r of toolResults) allToolCalls.push({ name: r.name, args: r.args, result: r.result });
154
184
 
155
185
  // Send function responses back to the model
156
186
  response = await this._withRetry(() => this.chatSession.sendMessage({
@@ -234,46 +264,90 @@ class ToolAgent extends BaseGemini {
234
264
  return;
235
265
  }
236
266
 
237
- // Execute tools sequentially so we can yield events
267
+ // Execute tools and yield events
238
268
  const toolResults = [];
239
- for (const call of functionCalls) {
240
- if (this._stopped) break;
269
+ if (this._concurrency === 1) {
270
+ // Sequential: yield tool_call, execute, yield tool_result for each
271
+ for (const call of functionCalls) {
272
+ if (this._stopped) break;
241
273
 
242
- yield { type: 'tool_call', toolName: call.name, args: call.args };
274
+ yield { type: 'tool_call', toolName: call.name, args: call.args };
243
275
 
244
- // Fire onToolCall callback
245
- if (this.onToolCall) {
246
- try { this.onToolCall(call.name, call.args); }
247
- catch (e) { log.warn(`onToolCall callback error: ${e.message}`); }
248
- }
276
+ if (this.onToolCall) {
277
+ try { this.onToolCall(call.name, call.args); }
278
+ catch (e) { log.warn(`onToolCall callback error: ${e.message}`); }
279
+ }
249
280
 
250
- // Check onBeforeExecution gate
251
- let denied = false;
252
- if (this.onBeforeExecution) {
253
- try {
254
- const allowed = await this.onBeforeExecution(call.name, call.args);
255
- if (allowed === false) denied = true;
256
- } catch (e) {
257
- log.warn(`onBeforeExecution callback error: ${e.message}`);
281
+ let denied = false;
282
+ if (this.onBeforeExecution) {
283
+ try {
284
+ const allowed = await this.onBeforeExecution(call.name, call.args);
285
+ if (allowed === false) denied = true;
286
+ } catch (e) {
287
+ log.warn(`onBeforeExecution callback error: ${e.message}`);
288
+ }
258
289
  }
259
- }
260
290
 
261
- let result;
262
- if (denied) {
263
- result = { error: 'Execution denied by onBeforeExecution callback' };
264
- } else {
265
- try {
266
- result = await this.toolExecutor(call.name, call.args);
267
- } catch (err) {
268
- log.warn(`Tool ${call.name} failed: ${err.message}`);
269
- result = { error: err.message };
291
+ let result;
292
+ if (denied) {
293
+ result = { error: 'Execution denied by onBeforeExecution callback' };
294
+ } else {
295
+ try {
296
+ result = await this.toolExecutor(call.name, call.args);
297
+ } catch (err) {
298
+ log.warn(`Tool ${call.name} failed: ${err.message}`);
299
+ result = { error: err.message };
300
+ }
270
301
  }
302
+
303
+ allToolCalls.push({ name: call.name, args: call.args, result });
304
+ yield { type: 'tool_result', toolName: call.name, result };
305
+
306
+ toolResults.push({ id: call.id, name: call.name, result });
307
+ }
308
+ } else {
309
+ // Parallel: yield all tool_call events, execute all, yield all tool_result events
310
+ for (const call of functionCalls) {
311
+ yield { type: 'tool_call', toolName: call.name, args: call.args };
271
312
  }
272
313
 
273
- allToolCalls.push({ name: call.name, args: call.args, result });
274
- yield { type: 'tool_result', toolName: call.name, result };
314
+ const tasks = functionCalls.map(call => async () => {
315
+ if (this.onToolCall) {
316
+ try { this.onToolCall(call.name, call.args); }
317
+ catch (e) { log.warn(`onToolCall callback error: ${e.message}`); }
318
+ }
275
319
 
276
- toolResults.push({ id: call.id, name: call.name, result });
320
+ let denied = false;
321
+ if (this.onBeforeExecution) {
322
+ try {
323
+ const allowed = await this.onBeforeExecution(call.name, call.args);
324
+ if (allowed === false) denied = true;
325
+ } catch (e) {
326
+ log.warn(`onBeforeExecution callback error: ${e.message}`);
327
+ }
328
+ }
329
+
330
+ let result;
331
+ if (denied) {
332
+ result = { error: 'Execution denied by onBeforeExecution callback' };
333
+ } else {
334
+ try {
335
+ result = await this.toolExecutor(call.name, call.args);
336
+ } catch (err) {
337
+ log.warn(`Tool ${call.name} failed: ${err.message}`);
338
+ result = { error: err.message };
339
+ }
340
+ }
341
+
342
+ return { id: call.id, name: call.name, args: call.args, result };
343
+ });
344
+
345
+ const results = await runWithConcurrency(tasks, this._concurrency);
346
+ for (const r of results) {
347
+ allToolCalls.push({ name: r.name, args: r.args, result: r.result });
348
+ yield { type: 'tool_result', toolName: r.name, result: r.result };
349
+ toolResults.push({ id: r.id, name: r.name, result: r.result });
350
+ }
277
351
  }
278
352
 
279
353
  // Send function responses back and get next stream
package/types.d.ts CHANGED
@@ -279,6 +279,8 @@ export interface ToolAgentOptions extends BaseGeminiOptions {
279
279
  onBeforeExecution?: (toolName: string, args: Record<string, any>) => Promise<boolean>;
280
280
  /** Directory for tool-written files (pass-through for toolExecutor use) */
281
281
  writeDir?: string;
282
+ /** Parallel tool execution: false = sequential, true = unlimited parallel, number = concurrency limit (default: true) */
283
+ parallelToolCalls?: boolean | number;
282
284
  }
283
285
 
284
286
  export interface LocalDataEntry {
@@ -539,6 +541,7 @@ export declare class ToolAgent extends BaseGemini {
539
541
  onBeforeExecution: ((toolName: string, args: Record<string, any>) => Promise<boolean>) | null;
540
542
  /** Directory for tool-written files (pass-through for toolExecutor use) */
541
543
  writeDir: string | null;
544
+ parallelToolCalls: boolean | number;
542
545
 
543
546
  chat(message: string, opts?: { labels?: Record<string, string> }): Promise<AgentResponse>;
544
547
  stream(message: string, opts?: { labels?: Record<string, string> }): AsyncGenerator<AgentStreamEvent, void, unknown>;