Ability to install missing packages

package.json

@@ -149,7 +149,9 @@
             },
             {
                 "name": "dachat_data_runPython",
-                "tags": [],
+                "tags": [
+                    "Python Execution"
+                ],
                 "displayName": "Run Python",
                 "modelDescription": "Execute Python code locally using Pyodide, providing access to Python's extensive functionality. This tool extends the LLM's capabilities by allowing it to run Python code for a wide range of computational tasks and data manipulations that it cannot perform directly. When you know the workspace folder path and the file path, use the relative path to the file when generating code.",
                 "inputSchema": {
@@ -168,6 +170,24 @@
                         "code"
                     ]
                 }
+            },
+            {
+                "name": "dachat_data_installPythonPackage",
+                "tags": [],
+                "displayName": "Install Missing Python Packages",
+                "modelDescription": "Install missing Python packages in the tool used to run Python code using Pyodide.",
+                "inputSchema": {
+                    "type": "object",
+                    "properties": {
+                        "package": {
+                            "type": "string",
+                            "description": "Name of the Python package that could not be found and needs to be installed"
+                        }
+                    },
+                    "required": [
+                        "code"
+                    ]
+                }
             }
         ]
     },

src/dataAgent.ts

@@ -96,7 +96,7 @@ export class DataAgent implements vscode.Disposable {
 				options.tools = allTools;
 			}
 
-			logger.debug('Sending request', JSON.stringify(messages, undefined, 4));
+			logger.debug('Sending request', JSON.stringify(messages));
 			const toolCalls: vscode.LanguageModelToolCallPart[] = [];
 
 			stream.progress('Analyzing');

src/extension.ts

@@ -6,7 +6,7 @@ import * as vscode from 'vscode';
 import { registerCsvCommand } from './csvCommand';
 import { DataAgent } from './dataAgent';
 import { initializeLogger } from './logger';
-import { FindFilesTool, RunPythonTool } from './tools';
+import { FindFilesTool, InstallPythonPackageTool, RunPythonTool } from './tools';
 
 export function activate(context: vscode.ExtensionContext) {
 	const dataAgent = new DataAgent(context);
@@ -15,7 +15,9 @@ export function activate(context: vscode.ExtensionContext) {
 	context.subscriptions.push(dataAgent);
 	context.subscriptions.push(registerCsvCommand());
 	context.subscriptions.push(vscode.lm.registerTool(FindFilesTool.Id, new FindFilesTool(context)));
-	context.subscriptions.push(vscode.lm.registerTool(RunPythonTool.Id, new RunPythonTool(context)));
+	const pythonTool = new RunPythonTool(context);
+	context.subscriptions.push(vscode.lm.registerTool(RunPythonTool.Id, pythonTool));
+	context.subscriptions.push(vscode.lm.registerTool(InstallPythonPackageTool.Id, new InstallPythonPackageTool(pythonTool)));
 
 	if (context.extensionMode === vscode.ExtensionMode.Test) {
 		return {

src/tools.ts

@@ -173,3 +173,42 @@ function sanitizePythonCode(code: string) {
 	}
 	return code;
 }
+
+
+interface IInstallPythonPackage {
+	package: string;
+}
+
+export class InstallPythonPackageTool implements vscode.LanguageModelTool<IInstallPythonPackage> {
+	public static Id = 'dachat_data_installPythonPackage';
+	constructor(readonly pythonTool: RunPythonTool) {
+	}
+
+	async invoke(
+		options: vscode.LanguageModelToolInvocationOptions<IInstallPythonPackage>,
+		token: vscode.CancellationToken
+	) {
+		logger.debug(`Installing Package "${options.input.package}"`);
+		const result = await this.pythonTool.invoke({
+			input: {
+				code: `import ${options.input.package}`,
+				reason: `Installing ${options.input.package}`
+			}, toolInvocationToken: options.toolInvocationToken,
+			tokenizationOptions: options.tokenizationOptions
+		}, token);
+
+		logger.debug(`Result after installing package ${options.input.package} => `);
+		logger.debug(JSON.stringify(result));
+
+		return new vscode.LanguageModelToolResult([new vscode.LanguageModelTextPart('Installation successful')]);
+	}
+
+	async prepareInvocation(
+		options: vscode.LanguageModelToolInvocationPrepareOptions<IInstallPythonPackage>,
+		_token: vscode.CancellationToken
+	) {
+		return {
+			invocationMessage: `Installing ${options.input.package}`
+		};
+	}
+}