ak-gemini 2.1.3 → 2.2.1

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';
 
 /**
@@ -19,7 +19,54 @@ import { randomUUID } from 'node:crypto';
 
 const MAX_OUTPUT_CHARS = 50_000;
 const MAX_FILE_TREE_LINES = 500;
-const IGNORE_DIRS = new Set(['node_modules', '.git', 'dist', 'coverage', '.next', 'build', '__pycache__']);
+const IGNORE_DIRS = new Set(['node_modules', '.git', 'dist', 'coverage', '.next', 'build', '__pycache__', '.venv']);
+
+const LANG_CONFIG = {
+	javascript: {
+		toolDescExecute: 'Execute a given piece of JavaScript code in a Node.js child process. Use this when you already have code to run — e.g., running code from a previous write_code call, re-running a snippet, or executing code the user provided. Use console.log() for output.',
+		toolDescWriteAndRun: 'Write a fresh solution from scratch and execute it in one step. Use this when you need to figure out the code AND run it — the autonomous, end-to-end tool for solving problems with code.',
+		codeParamDesc: 'JavaScript code to execute. Use console.log() for output. Use import syntax (ES modules).',
+		codeRules: `- Your code runs in a Node.js child process with access to all built-in modules
+- IMPORTANT: Your code runs as an ES module (.mjs). Use import syntax, NOT require():
+  - import fs from 'fs';
+  - import path from 'path';
+  - import { execSync } from 'child_process';
+- Use console.log() to produce output — that's how results are returned to you
+- Write efficient scripts that do multiple things per execution when possible
+- For parallel async operations, use Promise.all()
+- Handle errors in your scripts with try/catch so you get useful error messages
+- Top-level await is supported`,
+		commentsEnabled: `- Add a JSDoc @fileoverview comment at the top of each script explaining what it does\n- Add brief JSDoc @param comments for any functions you define`,
+		commentsDisabled: `- Do NOT write any comments in your code — save tokens. The code should be self-explanatory.`,
+		packageLabel: 'Available Packages',
+		packageIntro: 'These npm packages are installed and can be imported',
+		bashExamples: 'ls, grep, curl, git, npm, cat',
+		codeBlockLang: 'javascript',
+		execToolSummary: 'Run a given piece of JavaScript code. Use when you already have code to run — e.g., from a previous write_code call, re-running a snippet, or executing user-provided code.',
+		writeRunToolSummary: 'Write a fresh solution from scratch and execute it in one step. The autonomous, end-to-end tool for solving problems with code.',
+	},
+	python: {
+		toolDescExecute: 'Execute a given piece of Python code in a child process. Use this when you already have code to run — e.g., running code from a previous write_code call, re-running a snippet, or executing code the user provided. Use print() for output.',
+		toolDescWriteAndRun: 'Write a fresh Python solution from scratch and execute it in one step. Use this when you need to figure out the code AND run it — the autonomous, end-to-end tool for solving problems with code.',
+		codeParamDesc: 'Python code to execute. Use print() for output.',
+		codeRules: `- Your code runs in a Python 3 child process
+- Use print() to produce output — that's how results are returned to you
+- Use standard Python imports (import os, import json, from pathlib import Path, etc.)
+- Write efficient scripts; prefer list comprehensions and built-in functions
+- For async operations, use asyncio
+- Handle errors with try/except so you get useful error messages
+- A virtual environment is active — you can install packages using run_bash with: pip install <package>
+- Installed packages persist across executions in this session`,
+		commentsEnabled: `- Add a module-level docstring at the top of each script explaining what it does\n- Add brief docstrings for any functions you define`,
+		commentsDisabled: `- Do NOT write any comments in your code — save tokens. The code should be self-explanatory.`,
+		packageLabel: 'Available Packages',
+		packageIntro: 'These Python packages are available for import',
+		bashExamples: 'ls, grep, curl, git, pip, cat',
+		codeBlockLang: 'python',
+		execToolSummary: 'Run a given piece of Python code. Use when you already have code to run — e.g., from a previous write_code call, re-running a snippet, or executing user-provided code.',
+		writeRunToolSummary: 'Write a fresh Python solution from scratch and execute it in one step. The autonomous, end-to-end tool for solving problems with code.',
+	}
+};
 
 /** Tools that execute code/commands and can fail */
 const EXECUTING_TOOLS = new Set(['execute_code', 'write_and_run_code', 'run_bash']);
@@ -37,6 +84,8 @@ class CodeAgent extends BaseGemini {
 
 		// ── Agent config ──
 		this.workingDirectory = options.workingDirectory || process.cwd();
+		this.language = options.language || 'javascript';
+		this.pythonPath = options.pythonPath || null;
 		this.maxRounds = options.maxRounds || 10;
 		this.timeout = options.timeout || 30_000;
 		this.onBeforeExecution = options.onBeforeExecution || null;
@@ -49,6 +98,22 @@ class CodeAgent extends BaseGemini {
 		this.skills = options.skills || [];
 		this.envOverview = options.envOverview || '';
 
+		// ── Python state (resolved during init) ──
+		this._pythonBinary = null;
+		this._venvPath = null;
+		this._venvEnv = null;
+
+		// ── 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;
@@ -74,6 +139,8 @@ class CodeAgent extends BaseGemini {
 	 * @returns {{ functionDeclarations: Array<Object> }}
 	 */
 	_buildToolDefinitions() {
+		const lang = LANG_CONFIG[this.language] || LANG_CONFIG.javascript;
+
 		/** @type {Array<Object>} */
 		const declarations = [
 			{
@@ -91,11 +158,11 @@ class CodeAgent extends BaseGemini {
 			},
 			{
 				name: 'execute_code',
-				description: 'Execute a given piece of JavaScript code in a Node.js child process. Use this when you already have code to run — e.g., running code from a previous write_code call, re-running a snippet, or executing code the user provided. Use console.log() for output.',
+				description: lang.toolDescExecute,
 				parametersJsonSchema: {
 					type: 'object',
 					properties: {
-						code: { type: 'string', description: 'JavaScript code to execute. Use console.log() for output. Use import syntax (ES modules).' },
+						code: { type: 'string', description: lang.codeParamDesc },
 						purpose: { type: 'string', description: 'A short 2-4 word slug describing what this script does (e.g., "read-config", "parse-logs").' }
 					},
 					required: ['code']
@@ -103,11 +170,11 @@ class CodeAgent extends BaseGemini {
 			},
 			{
 				name: 'write_and_run_code',
-				description: 'Write a fresh solution from scratch and execute it in one step. Use this when you need to figure out the code AND run it — the autonomous, end-to-end tool for solving problems with code.',
+				description: lang.toolDescWriteAndRun,
 				parametersJsonSchema: {
 					type: 'object',
 					properties: {
-						code: { type: 'string', description: 'JavaScript code to write and execute. Use console.log() for output. Use import syntax (ES modules).' },
+						code: { type: 'string', description: lang.codeParamDesc },
 						purpose: { type: 'string', description: 'A short 2-4 word slug describing what this script does (e.g., "fetch-api-data", "generate-report").' }
 					},
 					required: ['code']
@@ -129,7 +196,7 @@ class CodeAgent extends BaseGemini {
 			},
 			{
 				name: 'run_bash',
-				description: 'Execute a shell command in the working directory. Use this for file operations, git commands, installing packages, or any shell task. Prefer this over execute_code for simple shell operations.',
+				description: `Execute a shell command in the working directory. Use this for file operations, git commands, installing packages, or any shell task (e.g., ${lang.bashExamples}). Prefer this over execute_code for simple shell operations.`,
 				parametersJsonSchema: {
 					type: 'object',
 					properties: {
@@ -156,6 +223,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 };
 	}
 
@@ -173,6 +245,12 @@ class CodeAgent extends BaseGemini {
 			await this._loadSkills();
 		}
 
+		// Resolve Python and set up venv if needed
+		if (this.language === 'python' && (!this._pythonBinary || force)) {
+			await this._resolvePython();
+			await this._setupVenv();
+		}
+
 		// Rebuild tools (use_skill may now be included)
 		this.chatConfig.tools = [this._buildToolDefinitions()];
 
@@ -211,6 +289,73 @@ class CodeAgent extends BaseGemini {
 		}
 	}
 
+	// ── Python Resolution ───────────────────────────────────────────────────
+
+	/**
+	 * Resolve the Python 3 binary path.
+	 * @private
+	 */
+	async _resolvePython() {
+		const tryBinary = (bin) => new Promise((resolve) => {
+			execFile(bin, ['--version'], { timeout: 5000 }, (err, stdout, stderr) => {
+				if (err) return resolve(null);
+				const output = (stdout || '') + (stderr || '');
+				if (output.includes('Python 3.')) return resolve(bin);
+				resolve(null);
+			});
+		});
+
+		if (this.pythonPath) {
+			const result = await tryBinary(this.pythonPath);
+			if (result) { this._pythonBinary = result; return; }
+			throw new Error(`CodeAgent: pythonPath "${this.pythonPath}" is not a valid Python 3 binary.`);
+		}
+
+		const python3 = await tryBinary('python3');
+		if (python3) { this._pythonBinary = python3; return; }
+
+		const python = await tryBinary('python');
+		if (python) { this._pythonBinary = python; return; }
+
+		throw new Error('CodeAgent: language is "python" but python3 was not found on PATH. Install Python 3 or provide the pythonPath option.');
+	}
+
+	/**
+	 * Create a virtual environment for Python execution.
+	 * @private
+	 */
+	async _setupVenv() {
+		await mkdir(this.writeDir, { recursive: true });
+		this._venvPath = join(this.writeDir, '.venv');
+		const isWin = process.platform === 'win32';
+		const venvBin = isWin ? join(this._venvPath, 'Scripts') : join(this._venvPath, 'bin');
+		const venvPython = join(venvBin, isWin ? 'python.exe' : 'python');
+
+		// Create venv if it doesn't exist
+		try {
+			await readFile(venvPython);
+		} catch {
+			log.debug(`Creating Python venv at ${this._venvPath}`);
+			await new Promise((resolve, reject) => {
+				execFile(this._pythonBinary, ['-m', 'venv', this._venvPath], {
+					timeout: 30_000
+				}, (err) => {
+					if (err) return reject(new Error(`CodeAgent: failed to create venv: ${err.message}`));
+					resolve();
+				});
+			});
+		}
+
+		// Build env with venv activated
+		const env = Object.assign({}, process.env);
+		env.VIRTUAL_ENV = this._venvPath;
+		env.PATH = venvBin + (isWin ? ';' : ':') + (process.env.PATH || '');
+		this._venvEnv = env;
+		this._pythonBinary = venvPython;
+
+		log.debug(`Python venv ready at ${this._venvPath}`);
+	}
+
 	// ── Context Gathering ────────────────────────────────────────────────────
 
 	/**
@@ -232,15 +377,37 @@ class CodeAgent extends BaseGemini {
 			fileTree = `${truncated}\n... (${lines.length - MAX_FILE_TREE_LINES} more files)`;
 		}
 
-		let npmPackages = [];
-		try {
-			const pkgPath = join(this.workingDirectory, 'package.json');
-			const pkg = JSON.parse(await readFile(pkgPath, 'utf-8'));
-			npmPackages = [
-				...Object.keys(pkg.dependencies || {}),
-				...Object.keys(pkg.devDependencies || {})
-			];
-		} catch { /* no package.json */ }
+		let packages = [];
+		if (this.language === 'python') {
+			try {
+				const reqPath = join(this.workingDirectory, 'requirements.txt');
+				const content = await readFile(reqPath, 'utf-8');
+				packages = content.split('\n')
+					.map(l => l.trim())
+					.filter(l => l && !l.startsWith('#') && !l.startsWith('-'))
+					.map(l => l.split(/[>=<!\[;\s]/)[0]);
+			} catch { /* no requirements.txt */ }
+			if (packages.length === 0) {
+				try {
+					const ppPath = join(this.workingDirectory, 'pyproject.toml');
+					const content = await readFile(ppPath, 'utf-8');
+					const depMatch = content.match(/dependencies\s*=\s*\[([\s\S]*?)\]/);
+					if (depMatch) {
+						packages = (depMatch[1].match(/"([^"]+)"/g) || [])
+							.map(s => s.replace(/"/g, '').split(/[>=<!\[;\s]/)[0]);
+					}
+				} catch { /* no pyproject.toml */ }
+			}
+		} else {
+			try {
+				const pkgPath = join(this.workingDirectory, 'package.json');
+				const pkg = JSON.parse(await readFile(pkgPath, 'utf-8'));
+				packages = [
+					...Object.keys(pkg.dependencies || {}),
+					...Object.keys(pkg.devDependencies || {})
+				];
+			} catch { /* no package.json */ }
+		}
 
 		const importantFileContents = [];
 		if (this.importantFiles.length > 0) {
@@ -252,7 +419,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) {
@@ -261,7 +428,7 @@ class CodeAgent extends BaseGemini {
 			}
 		}
 
-		this._codebaseContext = { fileTree, npmPackages, importantFileContents };
+		this._codebaseContext = { fileTree, npmPackages: packages, packages, importantFileContents };
 		this._contextGathered = true;
 	}
 
@@ -269,6 +436,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;
 
@@ -325,7 +494,8 @@ class CodeAgent extends BaseGemini {
 	 * @private
 	 */
 	_buildSystemPrompt() {
-		const { fileTree, npmPackages, importantFileContents } = this._codebaseContext || { fileTree: '', npmPackages: [], importantFileContents: [] };
+		const { fileTree, packages, importantFileContents } = this._codebaseContext || { fileTree: '', packages: [], importantFileContents: [] };
+		const lang = LANG_CONFIG[this.language] || LANG_CONFIG.javascript;
 
 		let prompt = `You are a coding agent working in ${this.workingDirectory}.
 
@@ -335,16 +505,16 @@ class CodeAgent extends BaseGemini {
 Output code without executing it. Use when showing, proposing, or presenting code to the user.
 
 ### execute_code
-Run a given piece of JavaScript code. Use when you already have code to run — e.g., from a previous write_code call, re-running a snippet, or executing user-provided code.
+${lang.execToolSummary}
 
 ### write_and_run_code
-Write a fresh solution from scratch and execute it in one step. The autonomous, end-to-end tool for solving problems with code.
+${lang.writeRunToolSummary}
 
 ### fix_code
 Fix broken code by providing original and fixed versions. Set execute=true to verify the fix works.
 
 ### run_bash
-Run shell commands directly (e.g., ls, grep, curl, git, npm, cat). Prefer this over execute_code for simple shell operations.`;
+Run shell commands directly (e.g., ${lang.bashExamples}). Prefer this over execute_code for simple shell operations.`;
 
 		if (this._skillRegistry.size > 0) {
 			prompt += `
@@ -358,36 +528,27 @@ Load a skill by name to get detailed instructions and templates. Available skill
 ## Code Execution Rules
 These rules apply when using execute_code, write_and_run_code, or fix_code (with execute=true):
 - Always provide a short descriptive \`purpose\` parameter (2-4 word slug like "read-config")
-- Your code runs in a Node.js child process with access to all built-in modules
-- IMPORTANT: Your code runs as an ES module (.mjs). Use import syntax, NOT require():
-  - import fs from 'fs';
-  - import path from 'path';
-  - import { execSync } from 'child_process';
-- Use console.log() to produce output — that's how results are returned to you
-- Write efficient scripts that do multiple things per execution when possible
-- For parallel async operations, use Promise.all()
-- Handle errors in your scripts with try/catch so you get useful error messages
-- Top-level await is supported
+${lang.codeRules}
 - The working directory is: ${this.workingDirectory}`;
 
 		if (this.comments) {
-			prompt += `\n- Add a JSDoc @fileoverview comment at the top of each script explaining what it does\n- Add brief JSDoc @param comments for any functions you define`;
+			prompt += `\n${lang.commentsEnabled}`;
 		} else {
-			prompt += `\n- Do NOT write any comments in your code — save tokens. The code should be self-explanatory.`;
+			prompt += `\n${lang.commentsDisabled}`;
 		}
 
 		if (fileTree) {
 			prompt += `\n\n## File Tree\n\`\`\`\n${fileTree}\n\`\`\``;
 		}
 
-		if (npmPackages.length > 0) {
-			prompt += `\n\n## Available Packages\nThese npm packages are installed and can be imported: ${npmPackages.join(', ')}`;
+		if (packages && packages.length > 0) {
+			prompt += `\n\n## ${lang.packageLabel}\n${lang.packageIntro}: ${packages.join(', ')}`;
 		}
 
 		if (importantFileContents && importantFileContents.length > 0) {
 			prompt += `\n\n## Key Files`;
 			for (const { path: filePath, content } of importantFileContents) {
-				prompt += `\n\n### ${filePath}\n\`\`\`javascript\n${content}\n\`\`\``;
+				prompt += `\n\n### ${filePath}\n\`\`\`${lang.codeBlockLang}\n${content}\n\`\`\``;
 			}
 		}
 
@@ -434,16 +595,20 @@ These rules apply when using execute_code, write_and_run_code, or fix_code (with
 		await mkdir(this.writeDir, { recursive: true });
 
 		const slug = this._slugify(purpose);
-		const tempFile = join(this.writeDir, `agent-${slug}-${Date.now()}.mjs`);
+		const ext = this.language === 'python' ? '.py' : '.mjs';
+		const tempFile = join(this.writeDir, `agent-${slug}-${Date.now()}${ext}`);
 
 		try {
 			await writeFile(tempFile, code, 'utf-8');
 
+			const binary = this.language === 'python' ? this._pythonBinary : 'node';
+			const execEnv = this.language === 'python' && this._venvEnv ? this._venvEnv : process.env;
+
 			const result = await new Promise((resolve) => {
-				const child = execFile('node', [tempFile], {
+				const child = execFile(binary, [tempFile], {
 					cwd: this.workingDirectory,
 					timeout: this.timeout,
-					env: process.env,
+					env: execEnv,
 					maxBuffer: 10 * 1024 * 1024
 				}, (err, stdout, stderr) => {
 					this._activeProcess = null;
@@ -513,11 +678,13 @@ These rules apply when using execute_code, write_and_run_code, or fix_code (with
 			}
 		}
 
+		const execEnv = this.language === 'python' && this._venvEnv ? this._venvEnv : process.env;
+
 		const result = await new Promise((resolve) => {
 			const child = execFile('bash', ['-c', command], {
 				cwd: this.workingDirectory,
 				timeout: this.timeout,
-				env: process.env,
+				env: execEnv,
 				maxBuffer: 10 * 1024 * 1024
 			}, (err, stdout, stderr) => {
 				this._activeProcess = null;
@@ -651,12 +818,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 +1044,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) {
@@ -922,8 +1112,9 @@ These rules apply when using execute_code, write_and_run_code, or fix_code (with
 	 * @returns {Array<{fileName: string, purpose: string|null, script: string, filePath: string|null, tool: string}>}
 	 */
 	dump() {
+		const ext = this.language === 'python' ? '.py' : '.mjs';
 		return this._allExecutions.map((exec, i) => ({
-			fileName: exec.purpose ? `agent-${exec.purpose}.mjs` : `script-${i + 1}.mjs`,
+			fileName: exec.purpose ? `agent-${exec.purpose}${ext}` : `script-${i + 1}${ext}`,
 			purpose: exec.purpose || null,
 			script: exec.code,
 			filePath: exec.filePath || null,

package/index.cjs

@@ -1626,7 +1626,55 @@ var import_node_path = require("node:path");
 var import_node_crypto = require("node:crypto");
 var MAX_OUTPUT_CHARS = 5e4;
 var MAX_FILE_TREE_LINES = 500;
-var IGNORE_DIRS = /* @__PURE__ */ new Set(["node_modules", ".git", "dist", "coverage", ".next", "build", "__pycache__"]);
+var IGNORE_DIRS = /* @__PURE__ */ new Set(["node_modules", ".git", "dist", "coverage", ".next", "build", "__pycache__", ".venv"]);
+var LANG_CONFIG = {
+  javascript: {
+    toolDescExecute: "Execute a given piece of JavaScript code in a Node.js child process. Use this when you already have code to run \u2014 e.g., running code from a previous write_code call, re-running a snippet, or executing code the user provided. Use console.log() for output.",
+    toolDescWriteAndRun: "Write a fresh solution from scratch and execute it in one step. Use this when you need to figure out the code AND run it \u2014 the autonomous, end-to-end tool for solving problems with code.",
+    codeParamDesc: "JavaScript code to execute. Use console.log() for output. Use import syntax (ES modules).",
+    codeRules: `- Your code runs in a Node.js child process with access to all built-in modules
+- IMPORTANT: Your code runs as an ES module (.mjs). Use import syntax, NOT require():
+  - import fs from 'fs';
+  - import path from 'path';
+  - import { execSync } from 'child_process';
+- Use console.log() to produce output \u2014 that's how results are returned to you
+- Write efficient scripts that do multiple things per execution when possible
+- For parallel async operations, use Promise.all()
+- Handle errors in your scripts with try/catch so you get useful error messages
+- Top-level await is supported`,
+    commentsEnabled: `- Add a JSDoc @fileoverview comment at the top of each script explaining what it does
+- Add brief JSDoc @param comments for any functions you define`,
+    commentsDisabled: `- Do NOT write any comments in your code \u2014 save tokens. The code should be self-explanatory.`,
+    packageLabel: "Available Packages",
+    packageIntro: "These npm packages are installed and can be imported",
+    bashExamples: "ls, grep, curl, git, npm, cat",
+    codeBlockLang: "javascript",
+    execToolSummary: "Run a given piece of JavaScript code. Use when you already have code to run \u2014 e.g., from a previous write_code call, re-running a snippet, or executing user-provided code.",
+    writeRunToolSummary: "Write a fresh solution from scratch and execute it in one step. The autonomous, end-to-end tool for solving problems with code."
+  },
+  python: {
+    toolDescExecute: "Execute a given piece of Python code in a child process. Use this when you already have code to run \u2014 e.g., running code from a previous write_code call, re-running a snippet, or executing code the user provided. Use print() for output.",
+    toolDescWriteAndRun: "Write a fresh Python solution from scratch and execute it in one step. Use this when you need to figure out the code AND run it \u2014 the autonomous, end-to-end tool for solving problems with code.",
+    codeParamDesc: "Python code to execute. Use print() for output.",
+    codeRules: `- Your code runs in a Python 3 child process
+- Use print() to produce output \u2014 that's how results are returned to you
+- Use standard Python imports (import os, import json, from pathlib import Path, etc.)
+- Write efficient scripts; prefer list comprehensions and built-in functions
+- For async operations, use asyncio
+- Handle errors with try/except so you get useful error messages
+- A virtual environment is active \u2014 you can install packages using run_bash with: pip install <package>
+- Installed packages persist across executions in this session`,
+    commentsEnabled: `- Add a module-level docstring at the top of each script explaining what it does
+- Add brief docstrings for any functions you define`,
+    commentsDisabled: `- Do NOT write any comments in your code \u2014 save tokens. The code should be self-explanatory.`,
+    packageLabel: "Available Packages",
+    packageIntro: "These Python packages are available for import",
+    bashExamples: "ls, grep, curl, git, pip, cat",
+    codeBlockLang: "python",
+    execToolSummary: "Run a given piece of Python code. Use when you already have code to run \u2014 e.g., from a previous write_code call, re-running a snippet, or executing user-provided code.",
+    writeRunToolSummary: "Write a fresh Python solution from scratch and execute it in one step. The autonomous, end-to-end tool for solving problems with code."
+  }
+};
 var EXECUTING_TOOLS = /* @__PURE__ */ new Set(["execute_code", "write_and_run_code", "run_bash"]);
 var CodeAgent = class extends base_default {
   /**
@@ -1638,6 +1686,8 @@ var CodeAgent = class extends base_default {
     }
     super(options);
     this.workingDirectory = options.workingDirectory || process.cwd();
+    this.language = options.language || "javascript";
+    this.pythonPath = options.pythonPath || null;
     this.maxRounds = options.maxRounds || 10;
     this.timeout = options.timeout || 3e4;
     this.onBeforeExecution = options.onBeforeExecution || null;
@@ -1649,6 +1699,18 @@ var CodeAgent = class extends base_default {
     this.maxRetries = options.maxRetries ?? 3;
     this.skills = options.skills || [];
     this.envOverview = options.envOverview || "";
+    this._pythonBinary = null;
+    this._venvPath = null;
+    this._venvEnv = null;
+    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;
@@ -1668,6 +1730,7 @@ var CodeAgent = class extends base_default {
    * @returns {{ functionDeclarations: Array<Object> }}
    */
   _buildToolDefinitions() {
+    const lang = LANG_CONFIG[this.language] || LANG_CONFIG.javascript;
     const declarations = [
       {
         name: "write_code",
@@ -1684,11 +1747,11 @@ var CodeAgent = class extends base_default {
       },
       {
         name: "execute_code",
-        description: "Execute a given piece of JavaScript code in a Node.js child process. Use this when you already have code to run \u2014 e.g., running code from a previous write_code call, re-running a snippet, or executing code the user provided. Use console.log() for output.",
+        description: lang.toolDescExecute,
         parametersJsonSchema: {
           type: "object",
           properties: {
-            code: { type: "string", description: "JavaScript code to execute. Use console.log() for output. Use import syntax (ES modules)." },
+            code: { type: "string", description: lang.codeParamDesc },
             purpose: { type: "string", description: 'A short 2-4 word slug describing what this script does (e.g., "read-config", "parse-logs").' }
           },
           required: ["code"]
@@ -1696,11 +1759,11 @@ var CodeAgent = class extends base_default {
       },
       {
         name: "write_and_run_code",
-        description: "Write a fresh solution from scratch and execute it in one step. Use this when you need to figure out the code AND run it \u2014 the autonomous, end-to-end tool for solving problems with code.",
+        description: lang.toolDescWriteAndRun,
         parametersJsonSchema: {
           type: "object",
           properties: {
-            code: { type: "string", description: "JavaScript code to write and execute. Use console.log() for output. Use import syntax (ES modules)." },
+            code: { type: "string", description: lang.codeParamDesc },
             purpose: { type: "string", description: 'A short 2-4 word slug describing what this script does (e.g., "fetch-api-data", "generate-report").' }
           },
           required: ["code"]
@@ -1722,7 +1785,7 @@ var CodeAgent = class extends base_default {
       },
       {
         name: "run_bash",
-        description: "Execute a shell command in the working directory. Use this for file operations, git commands, installing packages, or any shell task. Prefer this over execute_code for simple shell operations.",
+        description: `Execute a shell command in the working directory. Use this for file operations, git commands, installing packages, or any shell task (e.g., ${lang.bashExamples}). Prefer this over execute_code for simple shell operations.`,
         parametersJsonSchema: {
           type: "object",
           properties: {
@@ -1746,6 +1809,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 ─────────────────────────────────────────────────────────────────
@@ -1758,6 +1824,10 @@ var CodeAgent = class extends base_default {
     if (this.skills.length > 0 && (this._skillRegistry.size === 0 || force)) {
       await this._loadSkills();
     }
+    if (this.language === "python" && (!this._pythonBinary || force)) {
+      await this._resolvePython();
+      await this._setupVenv();
+    }
     this.chatConfig.tools = [this._buildToolDefinitions()];
     if (!this._contextGathered || force) {
       await this._gatherCodebaseContext();
@@ -1785,6 +1855,70 @@ var CodeAgent = class extends base_default {
       }
     }
   }
+  // ── Python Resolution ───────────────────────────────────────────────────
+  /**
+   * Resolve the Python 3 binary path.
+   * @private
+   */
+  async _resolvePython() {
+    const tryBinary = (bin) => new Promise((resolve2) => {
+      (0, import_node_child_process.execFile)(bin, ["--version"], { timeout: 5e3 }, (err, stdout, stderr) => {
+        if (err) return resolve2(null);
+        const output = (stdout || "") + (stderr || "");
+        if (output.includes("Python 3.")) return resolve2(bin);
+        resolve2(null);
+      });
+    });
+    if (this.pythonPath) {
+      const result = await tryBinary(this.pythonPath);
+      if (result) {
+        this._pythonBinary = result;
+        return;
+      }
+      throw new Error(`CodeAgent: pythonPath "${this.pythonPath}" is not a valid Python 3 binary.`);
+    }
+    const python3 = await tryBinary("python3");
+    if (python3) {
+      this._pythonBinary = python3;
+      return;
+    }
+    const python = await tryBinary("python");
+    if (python) {
+      this._pythonBinary = python;
+      return;
+    }
+    throw new Error('CodeAgent: language is "python" but python3 was not found on PATH. Install Python 3 or provide the pythonPath option.');
+  }
+  /**
+   * Create a virtual environment for Python execution.
+   * @private
+   */
+  async _setupVenv() {
+    await (0, import_promises2.mkdir)(this.writeDir, { recursive: true });
+    this._venvPath = (0, import_node_path.join)(this.writeDir, ".venv");
+    const isWin = process.platform === "win32";
+    const venvBin = isWin ? (0, import_node_path.join)(this._venvPath, "Scripts") : (0, import_node_path.join)(this._venvPath, "bin");
+    const venvPython = (0, import_node_path.join)(venvBin, isWin ? "python.exe" : "python");
+    try {
+      await (0, import_promises2.readFile)(venvPython);
+    } catch {
+      logger_default.debug(`Creating Python venv at ${this._venvPath}`);
+      await new Promise((resolve2, reject) => {
+        (0, import_node_child_process.execFile)(this._pythonBinary, ["-m", "venv", this._venvPath], {
+          timeout: 3e4
+        }, (err) => {
+          if (err) return reject(new Error(`CodeAgent: failed to create venv: ${err.message}`));
+          resolve2();
+        });
+      });
+    }
+    const env = Object.assign({}, process.env);
+    env.VIRTUAL_ENV = this._venvPath;
+    env.PATH = venvBin + (isWin ? ";" : ":") + (process.env.PATH || "");
+    this._venvEnv = env;
+    this._pythonBinary = venvPython;
+    logger_default.debug(`Python venv ready at ${this._venvPath}`);
+  }
   // ── Context Gathering ────────────────────────────────────────────────────
   /**
    * @private
@@ -1803,15 +1937,35 @@ var CodeAgent = class extends base_default {
       fileTree = `${truncated}
 ... (${lines.length - MAX_FILE_TREE_LINES} more files)`;
     }
-    let npmPackages = [];
-    try {
-      const pkgPath = (0, import_node_path.join)(this.workingDirectory, "package.json");
-      const pkg = JSON.parse(await (0, import_promises2.readFile)(pkgPath, "utf-8"));
-      npmPackages = [
-        ...Object.keys(pkg.dependencies || {}),
-        ...Object.keys(pkg.devDependencies || {})
-      ];
-    } catch {
+    let packages = [];
+    if (this.language === "python") {
+      try {
+        const reqPath = (0, import_node_path.join)(this.workingDirectory, "requirements.txt");
+        const content = await (0, import_promises2.readFile)(reqPath, "utf-8");
+        packages = content.split("\n").map((l) => l.trim()).filter((l) => l && !l.startsWith("#") && !l.startsWith("-")).map((l) => l.split(/[>=<!\[;\s]/)[0]);
+      } catch {
+      }
+      if (packages.length === 0) {
+        try {
+          const ppPath = (0, import_node_path.join)(this.workingDirectory, "pyproject.toml");
+          const content = await (0, import_promises2.readFile)(ppPath, "utf-8");
+          const depMatch = content.match(/dependencies\s*=\s*\[([\s\S]*?)\]/);
+          if (depMatch) {
+            packages = (depMatch[1].match(/"([^"]+)"/g) || []).map((s) => s.replace(/"/g, "").split(/[>=<!\[;\s]/)[0]);
+          }
+        } catch {
+        }
+      }
+    } else {
+      try {
+        const pkgPath = (0, import_node_path.join)(this.workingDirectory, "package.json");
+        const pkg = JSON.parse(await (0, import_promises2.readFile)(pkgPath, "utf-8"));
+        packages = [
+          ...Object.keys(pkg.dependencies || {}),
+          ...Object.keys(pkg.devDependencies || {})
+        ];
+      } catch {
+      }
     }
     const importantFileContents = [];
     if (this.importantFiles.length > 0) {
@@ -1823,7 +1977,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) {
@@ -1831,13 +1985,14 @@ var CodeAgent = class extends base_default {
         }
       }
     }
-    this._codebaseContext = { fileTree, npmPackages, importantFileContents };
+    this._codebaseContext = { fileTree, npmPackages: packages, packages, importantFileContents };
     this._contextGathered = true;
   }
   /**
    * @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(
@@ -1888,7 +2043,8 @@ var CodeAgent = class extends base_default {
    * @private
    */
   _buildSystemPrompt() {
-    const { fileTree, npmPackages, importantFileContents } = this._codebaseContext || { fileTree: "", npmPackages: [], importantFileContents: [] };
+    const { fileTree, packages, importantFileContents } = this._codebaseContext || { fileTree: "", packages: [], importantFileContents: [] };
+    const lang = LANG_CONFIG[this.language] || LANG_CONFIG.javascript;
     let prompt = `You are a coding agent working in ${this.workingDirectory}.
 
 ## Available Tools
@@ -1897,16 +2053,16 @@ var CodeAgent = class extends base_default {
 Output code without executing it. Use when showing, proposing, or presenting code to the user.
 
 ### execute_code
-Run a given piece of JavaScript code. Use when you already have code to run \u2014 e.g., from a previous write_code call, re-running a snippet, or executing user-provided code.
+${lang.execToolSummary}
 
 ### write_and_run_code
-Write a fresh solution from scratch and execute it in one step. The autonomous, end-to-end tool for solving problems with code.
+${lang.writeRunToolSummary}
 
 ### fix_code
 Fix broken code by providing original and fixed versions. Set execute=true to verify the fix works.
 
 ### run_bash
-Run shell commands directly (e.g., ls, grep, curl, git, npm, cat). Prefer this over execute_code for simple shell operations.`;
+Run shell commands directly (e.g., ${lang.bashExamples}). Prefer this over execute_code for simple shell operations.`;
     if (this._skillRegistry.size > 0) {
       prompt += `
 
@@ -1918,24 +2074,14 @@ Load a skill by name to get detailed instructions and templates. Available skill
 ## Code Execution Rules
 These rules apply when using execute_code, write_and_run_code, or fix_code (with execute=true):
 - Always provide a short descriptive \`purpose\` parameter (2-4 word slug like "read-config")
-- Your code runs in a Node.js child process with access to all built-in modules
-- IMPORTANT: Your code runs as an ES module (.mjs). Use import syntax, NOT require():
-  - import fs from 'fs';
-  - import path from 'path';
-  - import { execSync } from 'child_process';
-- Use console.log() to produce output \u2014 that's how results are returned to you
-- Write efficient scripts that do multiple things per execution when possible
-- For parallel async operations, use Promise.all()
-- Handle errors in your scripts with try/catch so you get useful error messages
-- Top-level await is supported
+${lang.codeRules}
 - The working directory is: ${this.workingDirectory}`;
     if (this.comments) {
       prompt += `
-- Add a JSDoc @fileoverview comment at the top of each script explaining what it does
-- Add brief JSDoc @param comments for any functions you define`;
+${lang.commentsEnabled}`;
     } else {
       prompt += `
-- Do NOT write any comments in your code \u2014 save tokens. The code should be self-explanatory.`;
+${lang.commentsDisabled}`;
     }
     if (fileTree) {
       prompt += `
@@ -1945,11 +2091,11 @@ These rules apply when using execute_code, write_and_run_code, or fix_code (with
 ${fileTree}
 \`\`\``;
     }
-    if (npmPackages.length > 0) {
+    if (packages && packages.length > 0) {
       prompt += `
 
-## Available Packages
-These npm packages are installed and can be imported: ${npmPackages.join(", ")}`;
+## ${lang.packageLabel}
+${lang.packageIntro}: ${packages.join(", ")}`;
     }
     if (importantFileContents && importantFileContents.length > 0) {
       prompt += `
@@ -1959,7 +2105,7 @@ These npm packages are installed and can be imported: ${npmPackages.join(", ")}`
         prompt += `
 
 ### ${filePath}
-\`\`\`javascript
+\`\`\`${lang.codeBlockLang}
 ${content}
 \`\`\``;
       }
@@ -2005,14 +2151,17 @@ ${this.envOverview}`;
     }
     await (0, import_promises2.mkdir)(this.writeDir, { recursive: true });
     const slug = this._slugify(purpose);
-    const tempFile = (0, import_node_path.join)(this.writeDir, `agent-${slug}-${Date.now()}.mjs`);
+    const ext = this.language === "python" ? ".py" : ".mjs";
+    const tempFile = (0, import_node_path.join)(this.writeDir, `agent-${slug}-${Date.now()}${ext}`);
     try {
       await (0, import_promises2.writeFile)(tempFile, code, "utf-8");
+      const binary = this.language === "python" ? this._pythonBinary : "node";
+      const execEnv = this.language === "python" && this._venvEnv ? this._venvEnv : process.env;
       const result = await new Promise((resolve2) => {
-        const child = (0, import_node_child_process.execFile)("node", [tempFile], {
+        const child = (0, import_node_child_process.execFile)(binary, [tempFile], {
           cwd: this.workingDirectory,
           timeout: this.timeout,
-          env: process.env,
+          env: execEnv,
           maxBuffer: 10 * 1024 * 1024
         }, (err, stdout, stderr) => {
           this._activeProcess = null;
@@ -2083,11 +2232,12 @@ ${this.envOverview}`;
         logger_default.warn(`onBeforeExecution callback error: ${e.message}`);
       }
     }
+    const execEnv = this.language === "python" && this._venvEnv ? this._venvEnv : process.env;
     const result = await new Promise((resolve2) => {
       const child = (0, import_node_child_process.execFile)("bash", ["-c", command], {
         cwd: this.workingDirectory,
         timeout: this.timeout,
-        env: process.env,
+        env: execEnv,
         maxBuffer: 10 * 1024 * 1024
       }, (err, stdout, stderr) => {
         this._activeProcess = null;
@@ -2234,12 +2384,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 ───────────────────────────────────────────────────
@@ -2396,6 +2564,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) {
@@ -2446,8 +2617,9 @@ ${this.envOverview}`;
    * @returns {Array<{fileName: string, purpose: string|null, script: string, filePath: string|null, tool: string}>}
    */
   dump() {
+    const ext = this.language === "python" ? ".py" : ".mjs";
     return this._allExecutions.map((exec, i) => ({
-      fileName: exec.purpose ? `agent-${exec.purpose}.mjs` : `script-${i + 1}.mjs`,
+      fileName: exec.purpose ? `agent-${exec.purpose}${ext}` : `script-${i + 1}${ext}`,
       purpose: exec.purpose || null,
       script: exec.code,
       filePath: exec.filePath || null,

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

package/types.d.ts

@@ -302,6 +302,10 @@ export interface RagAgentOptions extends BaseGeminiOptions {
 export interface CodeAgentOptions extends BaseGeminiOptions {
   /** Working directory for code execution (default: process.cwd()) */
   workingDirectory?: string;
+  /** Programming language for code execution: 'javascript' (default) or 'python' */
+  language?: 'javascript' | 'python';
+  /** Path to the Python binary (only used when language is 'python'; default: auto-detect python3/python) */
+  pythonPath?: string;
   /** Max code execution loop iterations (default: 10) */
   maxRounds?: number;
   /** Per-execution timeout in milliseconds (default: 30000) */
@@ -324,6 +328,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 +355,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 +385,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 +405,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 ──────────────────────────────────────────────────────
@@ -614,6 +637,8 @@ export declare class CodeAgent extends BaseGemini {
   constructor(options?: CodeAgentOptions);
 
   workingDirectory: string;
+  language: 'javascript' | 'python';
+  pythonPath: string | null;
   maxRounds: number;
   timeout: number;
   onBeforeExecution: ((content: string, toolName: string) => Promise<boolean> | boolean) | null;
@@ -625,6 +650,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>;