ak-gemini 2.0.6 → 2.0.9

package/GUIDE.md

@@ -349,6 +349,22 @@ console.log(result.text);       // "There were 47 new signups this week. I've se
 console.log(result.toolCalls);  // [{ name: 'query_db', args: {...}, result: [...] }, { name: 'send_email', ... }]
 ```
 
+### Parallel Tool Execution
+
+Control whether tool calls within a round execute in parallel or sequentially:
+
+```javascript
+const agent = new ToolAgent({
+  tools: [...],
+  toolExecutor: myExecutor,
+  parallelToolCalls: true,   // default: unlimited parallel execution
+  // parallelToolCalls: false, // sequential execution
+  // parallelToolCalls: 3,    // max 3 concurrent tool executions
+});
+```
+
+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.).
+
 ### Streaming
 
 Stream the agent's output in real-time — useful for showing progress in a UI:

package/README.md

@@ -413,6 +413,7 @@ All classes accept `BaseGeminiOptions`:
 | `maxToolRounds` | number | `10` | Max tool-use loop iterations |
 | `onToolCall` | function | — | Notification callback when tool is called |
 | `onBeforeExecution` | function | — | `async (toolName, args) => boolean` — gate execution |
+| `parallelToolCalls` | boolean \| number | `true` | Parallel tool execution: `false` = sequential, `true` = unlimited, number = concurrency limit |
 
 ### CodeAgent-Specific
 

package/base.js

@@ -98,6 +98,9 @@ class BaseGemini {
 		this.resourceExhaustedRetries = options.resourceExhaustedRetries ?? 5;
 		this.resourceExhaustedDelay = options.resourceExhaustedDelay ?? 1000;
 
+		// ── Health Check ──
+		this.healthCheck = options.healthCheck ?? false;
+
 		// ── Logging ──
 		this._configureLogLevel(options.logLevel);
 
@@ -192,11 +195,13 @@ class BaseGemini {
 		const chatOptions = this._getChatCreateOptions();
 		this.chatSession = this.genAIClient.chats.create(chatOptions);
 
-		try {
-			await this._withRetry(() => this.genAIClient.models.list());
-			log.debug(`${this.constructor.name}: API connection successful.`);
-		} catch (e) {
-			throw new Error(`${this.constructor.name} initialization failed: ${e.message}`);
+		if (this.healthCheck) {
+			try {
+				await this._withRetry(() => this.genAIClient.models.list());
+				log.debug(`${this.constructor.name}: API connection successful.`);
+			} catch (e) {
+				throw new Error(`${this.constructor.name} initialization failed: ${e.message}`);
+			}
 		}
 
 		log.debug(`${this.constructor.name}: Chat session initialized.`);

package/index.cjs

@@ -363,6 +363,7 @@ var BaseGemini = class {
     }
     this.resourceExhaustedRetries = options.resourceExhaustedRetries ?? 5;
     this.resourceExhaustedDelay = options.resourceExhaustedDelay ?? 1e3;
+    this.healthCheck = options.healthCheck ?? false;
     this._configureLogLevel(options.logLevel);
     this.labels = options.labels || {};
     this.enableGrounding = options.enableGrounding || false;
@@ -425,11 +426,13 @@ var BaseGemini = class {
     logger_default.debug(`Initializing ${this.constructor.name} chat session with model: ${this.modelName}...`);
     const chatOptions = this._getChatCreateOptions();
     this.chatSession = this.genAIClient.chats.create(chatOptions);
-    try {
-      await this._withRetry(() => this.genAIClient.models.list());
-      logger_default.debug(`${this.constructor.name}: API connection successful.`);
-    } catch (e) {
-      throw new Error(`${this.constructor.name} initialization failed: ${e.message}`);
+    if (this.healthCheck) {
+      try {
+        await this._withRetry(() => this.genAIClient.models.list());
+        logger_default.debug(`${this.constructor.name}: API connection successful.`);
+      } catch (e) {
+        throw new Error(`${this.constructor.name} initialization failed: ${e.message}`);
+      }
     }
     logger_default.debug(`${this.constructor.name}: Chat session initialized.`);
   }
@@ -1322,6 +1325,24 @@ var Message = class extends base_default {
 var message_default = Message;
 
 // tool-agent.js
+async function runWithConcurrency(tasks, concurrency) {
+  if (concurrency === Infinity) return Promise.all(tasks.map((t) => t()));
+  if (concurrency === 1) {
+    const results2 = [];
+    for (const t of tasks) results2.push(await t());
+    return results2;
+  }
+  const results = new Array(tasks.length);
+  let next = 0;
+  async function worker() {
+    while (next < tasks.length) {
+      const i = next++;
+      results[i] = await tasks[i]();
+    }
+  }
+  await Promise.all(Array.from({ length: Math.min(concurrency, tasks.length) }, () => worker()));
+  return results;
+}
 var ToolAgent = class extends base_default {
   /**
    * @param {ToolAgentOptions} [options={}]
@@ -1339,6 +1360,8 @@ var ToolAgent = class extends base_default {
     if (this.toolExecutor && this.tools.length === 0) {
       throw new Error("ToolAgent: toolExecutor provided without tools. Provide tool declarations so the model knows what tools are available.");
     }
+    this.parallelToolCalls = options.parallelToolCalls ?? true;
+    this._concurrency = this.parallelToolCalls === true ? Infinity : this.parallelToolCalls === false ? 1 : this.parallelToolCalls;
     this.maxToolRounds = options.maxToolRounds || 10;
     this.onToolCall = options.onToolCall || null;
     this.onBeforeExecution = options.onBeforeExecution || null;
@@ -1369,38 +1392,36 @@ var ToolAgent = class extends base_default {
       if (this._stopped) break;
       const functionCalls = response.functionCalls;
       if (!functionCalls || functionCalls.length === 0) break;
-      const toolResults = await Promise.all(
-        functionCalls.map(async (call) => {
-          if (this.onToolCall) {
-            try {
-              this.onToolCall(call.name, call.args);
-            } catch (e) {
-              logger_default.warn(`onToolCall callback error: ${e.message}`);
-            }
-          }
-          if (this.onBeforeExecution) {
-            try {
-              const allowed = await this.onBeforeExecution(call.name, call.args);
-              if (allowed === false) {
-                const result2 = { error: "Execution denied by onBeforeExecution callback" };
-                allToolCalls.push({ name: call.name, args: call.args, result: result2 });
-                return { id: call.id, name: call.name, result: result2 };
-              }
-            } catch (e) {
-              logger_default.warn(`onBeforeExecution callback error: ${e.message}`);
-            }
+      const tasks = functionCalls.map((call) => async () => {
+        if (this.onToolCall) {
+          try {
+            this.onToolCall(call.name, call.args);
+          } catch (e) {
+            logger_default.warn(`onToolCall callback error: ${e.message}`);
           }
-          let result;
+        }
+        if (this.onBeforeExecution) {
           try {
-            result = await this.toolExecutor(call.name, call.args);
-          } catch (err) {
-            logger_default.warn(`Tool ${call.name} failed: ${err.message}`);
-            result = { error: err.message };
+            const allowed = await this.onBeforeExecution(call.name, call.args);
+            if (allowed === false) {
+              const result2 = { error: "Execution denied by onBeforeExecution callback" };
+              return { id: call.id, name: call.name, args: call.args, result: result2 };
+            }
+          } catch (e) {
+            logger_default.warn(`onBeforeExecution callback error: ${e.message}`);
           }
-          allToolCalls.push({ name: call.name, args: call.args, result });
-          return { id: call.id, name: call.name, result };
-        })
-      );
+        }
+        let result;
+        try {
+          result = await this.toolExecutor(call.name, call.args);
+        } catch (err) {
+          logger_default.warn(`Tool ${call.name} failed: ${err.message}`);
+          result = { error: err.message };
+        }
+        return { id: call.id, name: call.name, args: call.args, result };
+      });
+      const toolResults = await runWithConcurrency(tasks, this._concurrency);
+      for (const r of toolResults) allToolCalls.push({ name: r.name, args: r.args, result: r.result });
       response = await this._withRetry(() => this.chatSession.sendMessage({
         message: toolResults.map((r) => ({
           functionResponse: {
@@ -1468,39 +1489,81 @@ var ToolAgent = class extends base_default {
         return;
       }
       const toolResults = [];
-      for (const call of functionCalls) {
-        if (this._stopped) break;
-        yield { type: "tool_call", toolName: call.name, args: call.args };
-        if (this.onToolCall) {
-          try {
-            this.onToolCall(call.name, call.args);
-          } catch (e) {
-            logger_default.warn(`onToolCall callback error: ${e.message}`);
+      if (this._concurrency === 1) {
+        for (const call of functionCalls) {
+          if (this._stopped) break;
+          yield { type: "tool_call", toolName: call.name, args: call.args };
+          if (this.onToolCall) {
+            try {
+              this.onToolCall(call.name, call.args);
+            } catch (e) {
+              logger_default.warn(`onToolCall callback error: ${e.message}`);
+            }
           }
-        }
-        let denied = false;
-        if (this.onBeforeExecution) {
-          try {
-            const allowed = await this.onBeforeExecution(call.name, call.args);
-            if (allowed === false) denied = true;
-          } catch (e) {
-            logger_default.warn(`onBeforeExecution callback error: ${e.message}`);
+          let denied = false;
+          if (this.onBeforeExecution) {
+            try {
+              const allowed = await this.onBeforeExecution(call.name, call.args);
+              if (allowed === false) denied = true;
+            } catch (e) {
+              logger_default.warn(`onBeforeExecution callback error: ${e.message}`);
+            }
           }
+          let result;
+          if (denied) {
+            result = { error: "Execution denied by onBeforeExecution callback" };
+          } else {
+            try {
+              result = await this.toolExecutor(call.name, call.args);
+            } catch (err) {
+              logger_default.warn(`Tool ${call.name} failed: ${err.message}`);
+              result = { error: err.message };
+            }
+          }
+          allToolCalls.push({ name: call.name, args: call.args, result });
+          yield { type: "tool_result", toolName: call.name, result };
+          toolResults.push({ id: call.id, name: call.name, result });
         }
-        let result;
-        if (denied) {
-          result = { error: "Execution denied by onBeforeExecution callback" };
-        } else {
-          try {
-            result = await this.toolExecutor(call.name, call.args);
-          } catch (err) {
-            logger_default.warn(`Tool ${call.name} failed: ${err.message}`);
-            result = { error: err.message };
+      } else {
+        for (const call of functionCalls) {
+          yield { type: "tool_call", toolName: call.name, args: call.args };
+        }
+        const tasks = functionCalls.map((call) => async () => {
+          if (this.onToolCall) {
+            try {
+              this.onToolCall(call.name, call.args);
+            } catch (e) {
+              logger_default.warn(`onToolCall callback error: ${e.message}`);
+            }
+          }
+          let denied = false;
+          if (this.onBeforeExecution) {
+            try {
+              const allowed = await this.onBeforeExecution(call.name, call.args);
+              if (allowed === false) denied = true;
+            } catch (e) {
+              logger_default.warn(`onBeforeExecution callback error: ${e.message}`);
+            }
           }
+          let result;
+          if (denied) {
+            result = { error: "Execution denied by onBeforeExecution callback" };
+          } else {
+            try {
+              result = await this.toolExecutor(call.name, call.args);
+            } catch (err) {
+              logger_default.warn(`Tool ${call.name} failed: ${err.message}`);
+              result = { error: err.message };
+            }
+          }
+          return { id: call.id, name: call.name, args: call.args, result };
+        });
+        const results = await runWithConcurrency(tasks, this._concurrency);
+        for (const r of results) {
+          allToolCalls.push({ name: r.name, args: r.args, result: r.result });
+          yield { type: "tool_result", toolName: r.name, result: r.result };
+          toolResults.push({ id: r.id, name: r.name, result: r.result });
         }
-        allToolCalls.push({ name: call.name, args: call.args, result });
-        yield { type: "tool_result", toolName: call.name, result };
-        toolResults.push({ id: call.id, name: call.name, result });
       }
       streamResponse = await this._withRetry(() => this.chatSession.sendMessageStream({
         message: toolResults.map((r) => ({

package/package.json

@@ -2,7 +2,7 @@
 	"name": "ak-gemini",
 	"author": "ak@mixpanel.com",
 	"description": "AK's Generative AI Helper for doing... everything",
-	"version": "2.0.6",
+	"version": "2.0.9",
 	"main": "index.js",
 	"files": [
 		"index.js",

package/tool-agent.js

@@ -13,6 +13,31 @@ import log from './logger.js';
  * @typedef {import('./types').AgentStreamEvent} AgentStreamEvent
  */
 
+/**
+ * Execute async task factories with a concurrency limit.
+ * @param {Array<() => Promise<any>>} tasks
+ * @param {number} concurrency - Infinity for unlimited, 1 for sequential
+ * @returns {Promise<any[]>} Results in same order as tasks
+ */
+async function runWithConcurrency(tasks, concurrency) {
+	if (concurrency === Infinity) return Promise.all(tasks.map(t => t()));
+	if (concurrency === 1) {
+		const results = [];
+		for (const t of tasks) results.push(await t());
+		return results;
+	}
+	const results = new Array(tasks.length);
+	let next = 0;
+	async function worker() {
+		while (next < tasks.length) {
+			const i = next++;
+			results[i] = await tasks[i]();
+		}
+	}
+	await Promise.all(Array.from({ length: Math.min(concurrency, tasks.length) }, () => worker()));
+	return results;
+}
+
 /**
  * AI agent that uses user-provided tools to accomplish tasks.
  * Automatically manages the tool-use loop: when the model decides to call
@@ -75,6 +100,13 @@ class ToolAgent extends BaseGemini {
 			throw new Error("ToolAgent: toolExecutor provided without tools. Provide tool declarations so the model knows what tools are available.");
 		}
 
+		// ── Parallel execution ──
+		this.parallelToolCalls = options.parallelToolCalls ?? true;
+		/** @private */
+		this._concurrency = this.parallelToolCalls === true ? Infinity
+			: this.parallelToolCalls === false ? 1
+			: this.parallelToolCalls;
+
 		// ── Tool loop config ──
 		this.maxToolRounds = options.maxToolRounds || 10;
 		this.onToolCall = options.onToolCall || null;
@@ -116,41 +148,39 @@ class ToolAgent extends BaseGemini {
 			const functionCalls = response.functionCalls;
 			if (!functionCalls || functionCalls.length === 0) break;
 
-			const toolResults = await Promise.all(
-				functionCalls.map(async (call) => {
-					// Fire onToolCall callback
-					if (this.onToolCall) {
-						try { this.onToolCall(call.name, call.args); }
-						catch (e) { log.warn(`onToolCall callback error: ${e.message}`); }
-					}
+			const tasks = functionCalls.map(call => async () => {
+				// Fire onToolCall callback
+				if (this.onToolCall) {
+					try { this.onToolCall(call.name, call.args); }
+					catch (e) { log.warn(`onToolCall callback error: ${e.message}`); }
+				}
 
-					// Check onBeforeExecution gate
-					if (this.onBeforeExecution) {
-						try {
-							const allowed = await this.onBeforeExecution(call.name, call.args);
-							if (allowed === false) {
-								const result = { error: 'Execution denied by onBeforeExecution callback' };
-								allToolCalls.push({ name: call.name, args: call.args, result });
-								return { id: call.id, name: call.name, result };
-							}
-						} catch (e) {
-							log.warn(`onBeforeExecution callback error: ${e.message}`);
+				// Check onBeforeExecution gate
+				if (this.onBeforeExecution) {
+					try {
+						const allowed = await this.onBeforeExecution(call.name, call.args);
+						if (allowed === false) {
+							const result = { error: 'Execution denied by onBeforeExecution callback' };
+							return { id: call.id, name: call.name, args: call.args, result };
 						}
+					} catch (e) {
+						log.warn(`onBeforeExecution callback error: ${e.message}`);
 					}
+				}
 
-					let result;
-					try {
-						result = await this.toolExecutor(call.name, call.args);
-					} catch (err) {
-						log.warn(`Tool ${call.name} failed: ${err.message}`);
-						result = { error: err.message };
-					}
+				let result;
+				try {
+					result = await this.toolExecutor(call.name, call.args);
+				} catch (err) {
+					log.warn(`Tool ${call.name} failed: ${err.message}`);
+					result = { error: err.message };
+				}
 
-					allToolCalls.push({ name: call.name, args: call.args, result });
+				return { id: call.id, name: call.name, args: call.args, result };
+			});
 
-					return { id: call.id, name: call.name, result };
-				})
-			);
+			const toolResults = await runWithConcurrency(tasks, this._concurrency);
+			for (const r of toolResults) allToolCalls.push({ name: r.name, args: r.args, result: r.result });
 
 			// Send function responses back to the model
 			response = await this._withRetry(() => this.chatSession.sendMessage({
@@ -234,46 +264,90 @@ class ToolAgent extends BaseGemini {
 				return;
 			}
 
-			// Execute tools sequentially so we can yield events
+			// Execute tools and yield events
 			const toolResults = [];
-			for (const call of functionCalls) {
-				if (this._stopped) break;
+			if (this._concurrency === 1) {
+				// Sequential: yield tool_call, execute, yield tool_result for each
+				for (const call of functionCalls) {
+					if (this._stopped) break;
 
-				yield { type: 'tool_call', toolName: call.name, args: call.args };
+					yield { type: 'tool_call', toolName: call.name, args: call.args };
 
-				// Fire onToolCall callback
-				if (this.onToolCall) {
-					try { this.onToolCall(call.name, call.args); }
-					catch (e) { log.warn(`onToolCall callback error: ${e.message}`); }
-				}
+					if (this.onToolCall) {
+						try { this.onToolCall(call.name, call.args); }
+						catch (e) { log.warn(`onToolCall callback error: ${e.message}`); }
+					}
 
-				// Check onBeforeExecution gate
-				let denied = false;
-				if (this.onBeforeExecution) {
-					try {
-						const allowed = await this.onBeforeExecution(call.name, call.args);
-						if (allowed === false) denied = true;
-					} catch (e) {
-						log.warn(`onBeforeExecution callback error: ${e.message}`);
+					let denied = false;
+					if (this.onBeforeExecution) {
+						try {
+							const allowed = await this.onBeforeExecution(call.name, call.args);
+							if (allowed === false) denied = true;
+						} catch (e) {
+							log.warn(`onBeforeExecution callback error: ${e.message}`);
+						}
 					}
-				}
 
-				let result;
-				if (denied) {
-					result = { error: 'Execution denied by onBeforeExecution callback' };
-				} else {
-					try {
-						result = await this.toolExecutor(call.name, call.args);
-					} catch (err) {
-						log.warn(`Tool ${call.name} failed: ${err.message}`);
-						result = { error: err.message };
+					let result;
+					if (denied) {
+						result = { error: 'Execution denied by onBeforeExecution callback' };
+					} else {
+						try {
+							result = await this.toolExecutor(call.name, call.args);
+						} catch (err) {
+							log.warn(`Tool ${call.name} failed: ${err.message}`);
+							result = { error: err.message };
+						}
 					}
+
+					allToolCalls.push({ name: call.name, args: call.args, result });
+					yield { type: 'tool_result', toolName: call.name, result };
+
+					toolResults.push({ id: call.id, name: call.name, result });
+				}
+			} else {
+				// Parallel: yield all tool_call events, execute all, yield all tool_result events
+				for (const call of functionCalls) {
+					yield { type: 'tool_call', toolName: call.name, args: call.args };
 				}
 
-				allToolCalls.push({ name: call.name, args: call.args, result });
-				yield { type: 'tool_result', toolName: call.name, result };
+				const tasks = functionCalls.map(call => async () => {
+					if (this.onToolCall) {
+						try { this.onToolCall(call.name, call.args); }
+						catch (e) { log.warn(`onToolCall callback error: ${e.message}`); }
+					}
 
-				toolResults.push({ id: call.id, name: call.name, result });
+					let denied = false;
+					if (this.onBeforeExecution) {
+						try {
+							const allowed = await this.onBeforeExecution(call.name, call.args);
+							if (allowed === false) denied = true;
+						} catch (e) {
+							log.warn(`onBeforeExecution callback error: ${e.message}`);
+						}
+					}
+
+					let result;
+					if (denied) {
+						result = { error: 'Execution denied by onBeforeExecution callback' };
+					} else {
+						try {
+							result = await this.toolExecutor(call.name, call.args);
+						} catch (err) {
+							log.warn(`Tool ${call.name} failed: ${err.message}`);
+							result = { error: err.message };
+						}
+					}
+
+					return { id: call.id, name: call.name, args: call.args, result };
+				});
+
+				const results = await runWithConcurrency(tasks, this._concurrency);
+				for (const r of results) {
+					allToolCalls.push({ name: r.name, args: r.args, result: r.result });
+					yield { type: 'tool_result', toolName: r.name, result: r.result };
+					toolResults.push({ id: r.id, name: r.name, result: r.result });
+				}
 			}
 
 			// Send function responses back and get next stream

package/types.d.ts

@@ -174,6 +174,9 @@ export interface BaseGeminiOptions {
   resourceExhaustedRetries?: number;
   /** Initial backoff delay in ms for 429 retries, doubles each attempt (default: 1000) */
   resourceExhaustedDelay?: number;
+
+  /** Run models.list() health check during init() (default: false) */
+  healthCheck?: boolean;
 }
 
 export interface TransformerOptions extends BaseGeminiOptions {
@@ -276,6 +279,8 @@ export interface ToolAgentOptions extends BaseGeminiOptions {
   onBeforeExecution?: (toolName: string, args: Record<string, any>) => Promise<boolean>;
   /** Directory for tool-written files (pass-through for toolExecutor use) */
   writeDir?: string;
+  /** Parallel tool execution: false = sequential, true = unlimited parallel, number = concurrency limit (default: true) */
+  parallelToolCalls?: boolean | number;
 }
 
 export interface LocalDataEntry {
@@ -536,6 +541,7 @@ export declare class ToolAgent extends BaseGemini {
   onBeforeExecution: ((toolName: string, args: Record<string, any>) => Promise<boolean>) | null;
   /** Directory for tool-written files (pass-through for toolExecutor use) */
   writeDir: string | null;
+  parallelToolCalls: boolean | number;
 
   chat(message: string, opts?: { labels?: Record<string, string> }): Promise<AgentResponse>;
   stream(message: string, opts?: { labels?: Record<string, string> }): AsyncGenerator<AgentStreamEvent, void, unknown>;