Initial code
This commit is contained in:
234
src/extension.ts
Normal file
234
src/extension.ts
Normal file
@@ -0,0 +1,234 @@
|
||||
// The module 'vscode' contains the VS Code extensibility API
|
||||
import * as vscode from 'vscode';
|
||||
import * as path from 'path';
|
||||
|
||||
// Custom TreeItem for files and folders
|
||||
class FileTreeItem extends vscode.TreeItem {
|
||||
constructor(
|
||||
public readonly resourceUri: vscode.Uri,
|
||||
public readonly collapsibleState: vscode.TreeItemCollapsibleState,
|
||||
public readonly isDirectory: boolean
|
||||
) {
|
||||
super(resourceUri, collapsibleState);
|
||||
this.contextValue = isDirectory ? 'directory' : 'file';
|
||||
this.tooltip = resourceUri.fsPath;
|
||||
this.description = isDirectory ? 'folder' : path.extname(resourceUri.fsPath);
|
||||
this.checkboxState = vscode.TreeItemCheckboxState.Unchecked;
|
||||
}
|
||||
}
|
||||
|
||||
// TreeView provider for the sidebar
|
||||
class PrompterTreeProvider implements vscode.TreeDataProvider<FileTreeItem> {
|
||||
private _onDidChangeTreeData: vscode.EventEmitter<FileTreeItem | undefined | null | void> = new vscode.EventEmitter<FileTreeItem | undefined | null | void>();
|
||||
readonly onDidChangeTreeData: vscode.Event<FileTreeItem | undefined | null | void> = this._onDidChangeTreeData.event;
|
||||
private selectedFiles: Set<string> = new Set();
|
||||
private xmlEditsEnabled: boolean = false;
|
||||
|
||||
constructor(private workspaceRoot: string | undefined) {}
|
||||
|
||||
getTreeItem(element: FileTreeItem): vscode.TreeItem {
|
||||
const item = element;
|
||||
if (this.selectedFiles.has(element.resourceUri.fsPath)) {
|
||||
item.checkboxState = vscode.TreeItemCheckboxState.Checked;
|
||||
item.label = path.basename(element.resourceUri.fsPath);
|
||||
} else {
|
||||
item.checkboxState = vscode.TreeItemCheckboxState.Unchecked;
|
||||
item.label = path.basename(element.resourceUri.fsPath);
|
||||
}
|
||||
return item;
|
||||
}
|
||||
|
||||
async getChildren(element?: FileTreeItem): Promise<FileTreeItem[]> {
|
||||
if (!this.workspaceRoot) {
|
||||
vscode.window.showInformationMessage('No workspace folder is opened');
|
||||
return Promise.resolve([]);
|
||||
}
|
||||
|
||||
if (element) {
|
||||
const dirPath = element.resourceUri.fsPath;
|
||||
return this.getFilesInDirectory(dirPath);
|
||||
} else {
|
||||
return this.getFilesInDirectory(this.workspaceRoot);
|
||||
}
|
||||
}
|
||||
|
||||
private async getFilesInDirectory(dirPath: string): Promise<FileTreeItem[]> {
|
||||
try {
|
||||
const files = await vscode.workspace.fs.readDirectory(vscode.Uri.file(dirPath));
|
||||
return Promise.all(files.map(async ([name, type]) => {
|
||||
const filePath = path.join(dirPath, name);
|
||||
const uri = vscode.Uri.file(filePath);
|
||||
const isDirectory = (type & vscode.FileType.Directory) !== 0;
|
||||
|
||||
return new FileTreeItem(
|
||||
uri,
|
||||
isDirectory ? vscode.TreeItemCollapsibleState.Collapsed : vscode.TreeItemCollapsibleState.None,
|
||||
isDirectory
|
||||
);
|
||||
}));
|
||||
} catch (error) {
|
||||
console.error(`Error reading directory ${dirPath}:`, error);
|
||||
return [];
|
||||
}
|
||||
}
|
||||
|
||||
toggleSelection(item: FileTreeItem): void {
|
||||
const filePath = item.resourceUri.fsPath;
|
||||
if (this.selectedFiles.has(filePath)) {
|
||||
this.selectedFiles.delete(filePath);
|
||||
} else {
|
||||
this.selectedFiles.add(filePath);
|
||||
}
|
||||
this._onDidChangeTreeData.fire();
|
||||
}
|
||||
|
||||
getSelectedFiles(): Set<string> {
|
||||
return this.selectedFiles;
|
||||
}
|
||||
|
||||
toggleXmlEdits(): void {
|
||||
this.xmlEditsEnabled = !this.xmlEditsEnabled;
|
||||
}
|
||||
|
||||
isXmlEditsEnabled(): boolean {
|
||||
return this.xmlEditsEnabled;
|
||||
}
|
||||
|
||||
refresh(): void {
|
||||
this._onDidChangeTreeData.fire();
|
||||
}
|
||||
}
|
||||
|
||||
// Utility function to estimate tokens in text
|
||||
function estimateTokens(text: string): number {
|
||||
// Rough estimation: Split by whitespace and punctuation
|
||||
// This is a simple approximation, actual token count may vary by model
|
||||
const words = text.split(/[\s\p{P}]+/u).filter(Boolean);
|
||||
return Math.ceil(words.length * 1.3); // Add 30% overhead for special tokens
|
||||
}
|
||||
|
||||
// Function to read file content
|
||||
async function readFileContent(filePath: string): Promise<string> {
|
||||
try {
|
||||
const readData = await vscode.workspace.fs.readFile(vscode.Uri.file(filePath));
|
||||
return Buffer.from(readData).toString('utf8');
|
||||
} catch (error) {
|
||||
console.error(`Error reading file ${filePath}:`, error);
|
||||
return '';
|
||||
}
|
||||
}
|
||||
|
||||
// Function to generate XML prompt
|
||||
function generateXMLPrompt(files: Map<string, { content: string; tokens: number }>, xmlEdits: boolean): string {
|
||||
const xmlParts = ['<?xml version="1.0" encoding="UTF-8"?>', '<prompt>'];
|
||||
|
||||
// Add files section
|
||||
xmlParts.push(' <files>');
|
||||
for (const [path, { content, tokens }] of files) {
|
||||
xmlParts.push(' <file>');
|
||||
xmlParts.push(` <path>${path}</path>`);
|
||||
xmlParts.push(` <tokens>${tokens}</tokens>`);
|
||||
xmlParts.push(` <content><![CDATA[${content}]]></content>`);
|
||||
xmlParts.push(' </file>');
|
||||
}
|
||||
xmlParts.push(' </files>');
|
||||
|
||||
// Add XML edits flag if enabled
|
||||
if (xmlEdits) {
|
||||
xmlParts.push(' <options>');
|
||||
xmlParts.push(' <xml_edits>true</xml_edits>');
|
||||
xmlParts.push(' </options>');
|
||||
}
|
||||
|
||||
xmlParts.push('</prompt>');
|
||||
return xmlParts.join('\n');
|
||||
}
|
||||
|
||||
// This method is called when your extension is activated
|
||||
export function activate(context: vscode.ExtensionContext) {
|
||||
console.log('Congratulations, your extension "prompter" is now active!');
|
||||
|
||||
const workspaceRoot = vscode.workspace.workspaceFolders?.[0]?.uri?.fsPath;
|
||||
const prompterTreeProvider = new PrompterTreeProvider(workspaceRoot);
|
||||
|
||||
// Register the TreeView
|
||||
const treeView = vscode.window.createTreeView('prompterView', {
|
||||
treeDataProvider: prompterTreeProvider,
|
||||
showCollapseAll: true
|
||||
});
|
||||
|
||||
// Create XML edits toggle button
|
||||
const xmlEditsButton = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left);
|
||||
xmlEditsButton.text = "$(diff-added) XML Edits";
|
||||
xmlEditsButton.tooltip = "Toggle XML Edits mode";
|
||||
xmlEditsButton.command = 'prompter.toggleXmlEdits';
|
||||
xmlEditsButton.show();
|
||||
|
||||
// Create copy button
|
||||
const copyButton = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Right);
|
||||
copyButton.text = "$(clippy) Copy";
|
||||
copyButton.tooltip = "Generate and copy XML prompt";
|
||||
copyButton.command = 'prompter.generatePrompt';
|
||||
copyButton.show();
|
||||
|
||||
// Register command to toggle file selection
|
||||
let toggleSelectionCommand = vscode.commands.registerCommand('prompter.toggleSelection', (item: FileTreeItem) => {
|
||||
prompterTreeProvider.toggleSelection(item);
|
||||
});
|
||||
|
||||
// Register command to toggle XML edits
|
||||
let toggleXmlEditsCommand = vscode.commands.registerCommand('prompter.toggleXmlEdits', () => {
|
||||
prompterTreeProvider.toggleXmlEdits();
|
||||
xmlEditsButton.text = prompterTreeProvider.isXmlEditsEnabled() ?
|
||||
"$(check) XML Edits" : "$(diff-added) XML Edits";
|
||||
});
|
||||
|
||||
// Register command to generate prompt from selected files
|
||||
let generatePromptCommand = vscode.commands.registerCommand('prompter.generatePrompt', async () => {
|
||||
const selectedFiles = prompterTreeProvider.getSelectedFiles();
|
||||
if (selectedFiles.size === 0) {
|
||||
vscode.window.showInformationMessage('Please select files first');
|
||||
return;
|
||||
}
|
||||
|
||||
let totalTokens = 0;
|
||||
const fileContents = new Map<string, { content: string; tokens: number }>();
|
||||
|
||||
// Process each selected file
|
||||
for (const filePath of selectedFiles) {
|
||||
const content = await readFileContent(filePath);
|
||||
const tokens = estimateTokens(content);
|
||||
totalTokens += tokens;
|
||||
fileContents.set(filePath, { content, tokens });
|
||||
|
||||
// Show individual file token count
|
||||
vscode.window.showInformationMessage(
|
||||
`File: ${filePath}\nEstimated tokens: ${tokens}`
|
||||
);
|
||||
}
|
||||
|
||||
// Show total token count
|
||||
vscode.window.showInformationMessage(
|
||||
`Total estimated tokens for ${selectedFiles.size} files: ${totalTokens}`
|
||||
);
|
||||
|
||||
// Generate XML prompt
|
||||
const xmlPrompt = generateXMLPrompt(fileContents, prompterTreeProvider.isXmlEditsEnabled());
|
||||
|
||||
// Copy to clipboard
|
||||
await vscode.env.clipboard.writeText(xmlPrompt);
|
||||
vscode.window.showInformationMessage('XML prompt copied to clipboard!');
|
||||
});
|
||||
|
||||
context.subscriptions.push(
|
||||
treeView,
|
||||
xmlEditsButton,
|
||||
copyButton,
|
||||
toggleSelectionCommand,
|
||||
toggleXmlEditsCommand,
|
||||
generatePromptCommand
|
||||
);
|
||||
}
|
||||
|
||||
// This method is called when your extension is deactivated
|
||||
export function deactivate() {}
|
||||
15
src/test/extension.test.ts
Normal file
15
src/test/extension.test.ts
Normal file
@@ -0,0 +1,15 @@
|
||||
import * as assert from 'assert';
|
||||
|
||||
// You can import and use all API from the 'vscode' module
|
||||
// as well as import your extension to test it
|
||||
import * as vscode from 'vscode';
|
||||
// import * as myExtension from '../../extension';
|
||||
|
||||
suite('Extension Test Suite', () => {
|
||||
vscode.window.showInformationMessage('Start all tests.');
|
||||
|
||||
test('Sample test', () => {
|
||||
assert.strictEqual(-1, [1, 2, 3].indexOf(5));
|
||||
assert.strictEqual(-1, [1, 2, 3].indexOf(0));
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user