ak-gemini 2.1.1 → 2.1.5

package/chat.js

@@ -9,6 +9,7 @@ import log from './logger.js';
 /**
  * @typedef {import('./types').ChatOptions} ChatOptions
  * @typedef {import('./types').ChatResponse} ChatResponse
+ * @typedef {import('./types').ChatStreamEvent} ChatStreamEvent
  */
 
 /**
@@ -82,6 +83,34 @@ class Chat extends BaseGemini {
 			usage: this.getLastUsage()
 		};
 	}
+
+	/**
+	 * Send a message and stream the response as events.
+	 *
+	 * @param {string} message - The user's message
+	 * @param {Object} [opts={}] - Per-message options
+	 * @yields {ChatStreamEvent}
+	 */
+	async *stream(message, opts = {}) {
+		if (!this.chatSession) await this.init();
+
+		let fullText = '';
+		const streamResponse = await this._withRetry(() => this.chatSession.sendMessageStream({ message }));
+
+		for await (const chunk of streamResponse) {
+			if (chunk.candidates?.[0]?.content?.parts?.[0]?.text) {
+				const text = chunk.candidates[0].content.parts[0].text;
+				fullText += text;
+				yield { type: 'text', text };
+			}
+		}
+
+		yield {
+			type: 'done',
+			fullText,
+			usage: this.getLastUsage()
+		};
+	}
 }
 
 export default Chat;

package/code-agent.js

@@ -8,7 +8,7 @@ import BaseGemini from './base.js';
 import log from './logger.js';
 import { execFile } from 'node:child_process';
 import { writeFile, unlink, readdir, readFile, mkdir } from 'node:fs/promises';
-import { join, sep, basename } from 'node:path';
+import { join, sep, basename, isAbsolute } from 'node:path';
 import { randomUUID } from 'node:crypto';
 
 /**
@@ -49,6 +49,17 @@ class CodeAgent extends BaseGemini {
 		this.skills = options.skills || [];
 		this.envOverview = options.envOverview || '';
 
+		// ── Custom tools ──
+		this.customTools = (options.tools || []).map(t => ({
+			name: t.name,
+			description: t.description,
+			parametersJsonSchema: t.parametersJsonSchema || t.parameters || t.input_schema || t.inputSchema
+		}));
+		this.toolExecutor = options.toolExecutor || null;
+		if (this.customTools.length > 0 && !this.toolExecutor) {
+			throw new Error('CodeAgent: tools provided without a toolExecutor.');
+		}
+
 		// ── Internal state ──
 		this._codebaseContext = null;
 		this._contextGathered = false;
@@ -156,6 +167,11 @@ class CodeAgent extends BaseGemini {
 			});
 		}
 
+		// Append custom tools
+		for (const t of this.customTools) {
+			declarations.push({ name: t.name, description: t.description, parametersJsonSchema: t.parametersJsonSchema });
+		}
+
 		return { functionDeclarations: declarations };
 	}
 
@@ -252,7 +268,7 @@ class CodeAgent extends BaseGemini {
 					continue;
 				}
 				try {
-					const fullPath = join(this.workingDirectory, resolved);
+					const fullPath = isAbsolute(resolved) ? resolved : join(this.workingDirectory, resolved);
 					const content = await readFile(fullPath, 'utf-8');
 					importantFileContents.push({ path: resolved, content });
 				} catch (e) {
@@ -269,6 +285,8 @@ class CodeAgent extends BaseGemini {
 	 * @private
 	 */
 	_resolveImportantFile(filename, fileTreeLines) {
+		if (isAbsolute(filename)) return filename;
+
 		const exact = fileTreeLines.find(line => line === filename);
 		if (exact) return exact;
 
@@ -651,12 +669,30 @@ These rules apply when using execute_code, write_and_run_code, or fix_code (with
 					data: { tool: 'use_skill', skillName: skill.name, content: skill.content, found: true }
 				};
 			}
-			default:
+			default: {
+				if (this.toolExecutor) {
+					try {
+						const result = await this.toolExecutor(name, input);
+						const resultStr = typeof result === 'string' ? result : JSON.stringify(result);
+						return {
+							output: resultStr,
+							type: 'tool',
+							data: { tool: name, args: input, result }
+						};
+					} catch (err) {
+						return {
+							output: `Tool "${name}" failed: ${err.message}`,
+							type: 'tool',
+							data: { tool: name, args: input, error: err.message }
+						};
+					}
+				}
 				return {
 					output: `Unknown tool: ${name}`,
 					type: 'unknown',
 					data: { tool: name }
 				};
+			}
 		}
 	}
 
@@ -859,6 +895,11 @@ These rules apply when using execute_code, write_and_run_code, or fix_code (with
 					yield { type: 'skill', skillName: data.skillName, content: data.content, found: data.found };
 				}
 
+				// Emit custom tool event
+				if (type === 'tool') {
+					yield { type: 'tool', toolName, args: data.args, result: data.result, error: data.error };
+				}
+
 				// Track consecutive failures
 				const isExecutingTool = EXECUTING_TOOLS.has(toolName) || (toolName === 'fix_code' && toolInput.execute);
 				if (isExecutingTool) {

package/index.cjs

@@ -1216,6 +1216,30 @@ var Chat = class extends base_default {
       usage: this.getLastUsage()
     };
   }
+  /**
+   * Send a message and stream the response as events.
+   *
+   * @param {string} message - The user's message
+   * @param {Object} [opts={}] - Per-message options
+   * @yields {ChatStreamEvent}
+   */
+  async *stream(message, opts = {}) {
+    if (!this.chatSession) await this.init();
+    let fullText = "";
+    const streamResponse = await this._withRetry(() => this.chatSession.sendMessageStream({ message }));
+    for await (const chunk of streamResponse) {
+      if (chunk.candidates?.[0]?.content?.parts?.[0]?.text) {
+        const text = chunk.candidates[0].content.parts[0].text;
+        fullText += text;
+        yield { type: "text", text };
+      }
+    }
+    yield {
+      type: "done",
+      fullText,
+      usage: this.getLastUsage()
+    };
+  }
 };
 var chat_default = Chat;
 
@@ -1625,6 +1649,15 @@ var CodeAgent = class extends base_default {
     this.maxRetries = options.maxRetries ?? 3;
     this.skills = options.skills || [];
     this.envOverview = options.envOverview || "";
+    this.customTools = (options.tools || []).map((t) => ({
+      name: t.name,
+      description: t.description,
+      parametersJsonSchema: t.parametersJsonSchema || t.parameters || t.input_schema || t.inputSchema
+    }));
+    this.toolExecutor = options.toolExecutor || null;
+    if (this.customTools.length > 0 && !this.toolExecutor) {
+      throw new Error("CodeAgent: tools provided without a toolExecutor.");
+    }
     this._codebaseContext = null;
     this._contextGathered = false;
     this._stopped = false;
@@ -1722,6 +1755,9 @@ var CodeAgent = class extends base_default {
         }
       });
     }
+    for (const t of this.customTools) {
+      declarations.push({ name: t.name, description: t.description, parametersJsonSchema: t.parametersJsonSchema });
+    }
     return { functionDeclarations: declarations };
   }
   // ── Init ─────────────────────────────────────────────────────────────────
@@ -1799,7 +1835,7 @@ var CodeAgent = class extends base_default {
           continue;
         }
         try {
-          const fullPath = (0, import_node_path.join)(this.workingDirectory, resolved);
+          const fullPath = (0, import_node_path.isAbsolute)(resolved) ? resolved : (0, import_node_path.join)(this.workingDirectory, resolved);
           const content = await (0, import_promises2.readFile)(fullPath, "utf-8");
           importantFileContents.push({ path: resolved, content });
         } catch (e) {
@@ -1814,6 +1850,7 @@ var CodeAgent = class extends base_default {
    * @private
    */
   _resolveImportantFile(filename, fileTreeLines) {
+    if ((0, import_node_path.isAbsolute)(filename)) return filename;
     const exact = fileTreeLines.find((line) => line === filename);
     if (exact) return exact;
     const partial = fileTreeLines.find(
@@ -2210,12 +2247,30 @@ ${this.envOverview}`;
           data: { tool: "use_skill", skillName: skill.name, content: skill.content, found: true }
         };
       }
-      default:
+      default: {
+        if (this.toolExecutor) {
+          try {
+            const result = await this.toolExecutor(name, input);
+            const resultStr = typeof result === "string" ? result : JSON.stringify(result);
+            return {
+              output: resultStr,
+              type: "tool",
+              data: { tool: name, args: input, result }
+            };
+          } catch (err) {
+            return {
+              output: `Tool "${name}" failed: ${err.message}`,
+              type: "tool",
+              data: { tool: name, args: input, error: err.message }
+            };
+          }
+        }
         return {
           output: `Unknown tool: ${name}`,
           type: "unknown",
           data: { tool: name }
         };
+      }
     }
   }
   // ── Non-Streaming Chat ───────────────────────────────────────────────────
@@ -2372,6 +2427,9 @@ ${this.envOverview}`;
         if (toolName === "use_skill") {
           yield { type: "skill", skillName: data.skillName, content: data.content, found: data.found };
         }
+        if (type === "tool") {
+          yield { type: "tool", toolName, args: data.args, result: data.result, error: data.error };
+        }
         const isExecutingTool = EXECUTING_TOOLS.has(toolName) || toolName === "fix_code" && toolInput.execute;
         if (isExecutingTool) {
           if (data.exitCode !== 0 && !data.denied) {

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.1.1",
+	"version": "2.1.5",
 	"main": "index.js",
 	"files": [
 		"index.js",

package/types.d.ts

@@ -324,6 +324,17 @@ export interface CodeAgentOptions extends BaseGeminiOptions {
   skills?: string[];
   /** Plain text environment overview appended to the system prompt — describe the project, stack, conventions, etc. */
   envOverview?: string;
+  /** Custom tool declarations to add alongside built-in CodeAgent tools. Accepts Gemini, Claude, or OpenAI tool formats (auto-mapped). */
+  tools?: Array<{
+    name: string;
+    description: string;
+    parametersJsonSchema?: any;
+    parameters?: any;
+    input_schema?: any;
+    inputSchema?: any;
+  }>;
+  /** Function to execute custom tool calls: (toolName, args) => result */
+  toolExecutor?: (toolName: string, args: Record<string, any>) => Promise<any>;
 }
 
 export interface CodeExecution {
@@ -340,7 +351,7 @@ export interface CodeExecution {
 }
 
 export interface ToolCallResult {
-  tool: 'write_code' | 'execute_code' | 'write_and_run_code' | 'fix_code' | 'run_bash' | 'use_skill';
+  tool: 'write_code' | 'execute_code' | 'write_and_run_code' | 'fix_code' | 'run_bash' | 'use_skill' | string;
   code?: string;
   purpose?: string;
   language?: string;
@@ -370,7 +381,7 @@ export interface CodeAgentResponse {
 }
 
 export interface CodeAgentStreamEvent {
-  type: 'text' | 'code' | 'output' | 'write' | 'fix' | 'bash' | 'skill' | 'done';
+  type: 'text' | 'code' | 'output' | 'write' | 'fix' | 'bash' | 'skill' | 'tool' | 'done';
   text?: string;
   code?: string;
   stdout?: string;
@@ -390,6 +401,14 @@ export interface CodeAgentStreamEvent {
   skillName?: string;
   content?: string;
   found?: boolean;
+  /** custom tool: tool name */
+  toolName?: string;
+  /** custom tool: arguments passed */
+  args?: Record<string, any>;
+  /** custom tool: result returned */
+  result?: any;
+  /** custom tool: error message (if failed) */
+  error?: string;
 }
 
 // ── Per-Message Options ──────────────────────────────────────────────────────
@@ -421,6 +440,13 @@ export interface ChatResponse {
   usage: UsageData | null;
 }
 
+export interface ChatStreamEvent {
+  type: 'text' | 'done';
+  text?: string;
+  fullText?: string;
+  usage?: UsageData | null;
+}
+
 export interface MessageResponse {
   /** The model's text response */
   text: string;
@@ -548,6 +574,7 @@ export declare class Chat extends BaseGemini {
   constructor(options?: ChatOptions);
 
   send(message: string, opts?: { labels?: Record<string, string> }): Promise<ChatResponse>;
+  stream(message: string, opts?: { labels?: Record<string, string> }): AsyncGenerator<ChatStreamEvent, void, unknown>;
 }
 
 export declare class Message extends BaseGemini {
@@ -617,6 +644,8 @@ export declare class CodeAgent extends BaseGemini {
   maxRetries: number;
   skills: string[];
   envOverview: string;
+  customTools: Array<{ name: string; description: string; parametersJsonSchema: any }>;
+  toolExecutor: ((toolName: string, args: Record<string, any>) => Promise<any>) | null;
 
   init(force?: boolean): Promise<void>;
   chat(message: string, opts?: { labels?: Record<string, string> }): Promise<CodeAgentResponse>;