292 lines
11 KiB
TypeScript
292 lines
11 KiB
TypeScript
import * as vscode from 'vscode';
|
|
import * as fs from 'fs';
|
|
import * as path from 'path';
|
|
import { PrompterSettings } from '../models/settings';
|
|
import { FileReader } from './fileReader';
|
|
import { TokenEstimator } from './tokenEstimator';
|
|
|
|
// Use require for ignore package with proper fallback
|
|
let ignoreFunc: () => any;
|
|
try {
|
|
const ignoreModule = require('ignore');
|
|
if (typeof ignoreModule === 'function') {
|
|
ignoreFunc = ignoreModule;
|
|
} else if (ignoreModule && typeof ignoreModule.default === 'function') {
|
|
ignoreFunc = ignoreModule.default;
|
|
} else {
|
|
throw new Error('Ignore module is neither a function nor has a default function');
|
|
}
|
|
console.log('Successfully loaded ignore function');
|
|
} catch (error) {
|
|
console.error('Error loading ignore package:', error);
|
|
ignoreFunc = () => ({
|
|
add: () => {},
|
|
ignores: () => false
|
|
});
|
|
}
|
|
|
|
export class PromptGenerator {
|
|
/**
|
|
* Generate a prompt from the selected files
|
|
* @param selectedFiles Set of file paths to include in the prompt
|
|
* @param settings Settings to apply when generating the prompt
|
|
* @returns The generated prompt text
|
|
*/
|
|
static async generatePrompt(selectedFiles: Set<string>, settings: PrompterSettings): Promise<string | null> {
|
|
if (selectedFiles.size === 0) {
|
|
throw new Error('No files selected');
|
|
}
|
|
|
|
const workspaceRoot = vscode.workspace.workspaceFolders?.[0]?.uri?.fsPath || '';
|
|
|
|
// Create ignore instance using the resolved function
|
|
const ig = ignoreFunc();
|
|
console.log('Created ignore instance');
|
|
|
|
// Load .gitignore patterns
|
|
try {
|
|
const gitignorePath = path.join(workspaceRoot, '.gitignore');
|
|
if (fs.existsSync(gitignorePath)) {
|
|
const gitignoreContent = fs.readFileSync(gitignorePath, 'utf8');
|
|
ig.add(gitignoreContent);
|
|
console.log('Successfully loaded .gitignore patterns:', gitignoreContent.split('\n').filter(Boolean));
|
|
} else {
|
|
console.log('No .gitignore file found at:', gitignorePath);
|
|
}
|
|
} catch (error) {
|
|
console.error('Error loading .gitignore:', error);
|
|
}
|
|
|
|
const fileContents = new Map<string, { content: string; tokens: number }>();
|
|
const filteredFiles = new Set<string>();
|
|
let totalTokens = 0;
|
|
|
|
// Process and filter files
|
|
for (const filePath of selectedFiles) {
|
|
try {
|
|
const stat = fs.statSync(filePath);
|
|
if (stat.isDirectory()) {
|
|
console.log(`Skipping directory: ${filePath}`);
|
|
continue;
|
|
}
|
|
} catch (error) {
|
|
console.error(`Error checking file stats for ${filePath}:`, error);
|
|
continue;
|
|
}
|
|
|
|
const relativePath = this.getRelativePath(filePath, workspaceRoot);
|
|
const normalizedPath = relativePath.split(path.sep).join('/');
|
|
|
|
console.log(`Checking path: ${normalizedPath}`);
|
|
if (ig.ignores(normalizedPath)) {
|
|
console.log(`Ignoring file based on patterns: ${normalizedPath}`);
|
|
continue;
|
|
}
|
|
|
|
console.log(`Processing file: ${normalizedPath}`);
|
|
filteredFiles.add(filePath);
|
|
const content = await FileReader.readFileContent(filePath);
|
|
const tokens = TokenEstimator.estimateTokens(content);
|
|
totalTokens += tokens;
|
|
fileContents.set(filePath, { content, tokens });
|
|
}
|
|
|
|
if (filteredFiles.size === 0) {
|
|
// Return null to signal that no files were available after filtering
|
|
return null;
|
|
}
|
|
|
|
// Create filtered contents map
|
|
const filteredContents = new Map<string, { content: string; tokens: number }>();
|
|
for (const filePath of filteredFiles) {
|
|
if (fileContents.has(filePath)) {
|
|
filteredContents.set(filePath, fileContents.get(filePath)!);
|
|
}
|
|
}
|
|
|
|
return this.generateXMLPrompt(filteredContents, settings);
|
|
}
|
|
|
|
|
|
/**
|
|
* Generate an XML formatted prompt following the new schema format
|
|
* @param files Map of file paths to content and token counts
|
|
* @param settings Settings to apply
|
|
* @returns XML formatted prompt
|
|
*/
|
|
private static generateXMLPrompt(files: Map<string, { content: string; tokens: number }>, settings: PrompterSettings): string {
|
|
const xmlParts: string[] = [];
|
|
|
|
let formattingInstructions = '';
|
|
if (settings.includeFormattingInstructions) {
|
|
try {
|
|
const extensionPath = vscode.extensions.getExtension('prompter')?.extensionPath ||
|
|
path.join(__dirname, '..', '..');
|
|
const formattingInstructionsPath = path.join(extensionPath, 'resources', 'xml_formatting_instructions.xml');
|
|
|
|
if (fs.existsSync(formattingInstructionsPath)) {
|
|
formattingInstructions = fs.readFileSync(formattingInstructionsPath, 'utf8');
|
|
} else {
|
|
console.warn('XML formatting instructions file not found at:', formattingInstructionsPath);
|
|
}
|
|
} catch (error) {
|
|
console.error('Error reading XML formatting instructions:', error);
|
|
}
|
|
}
|
|
|
|
if (settings.includeFileMap) {
|
|
xmlParts.push('<file_map>');
|
|
const workspaceRoot = vscode.workspace.workspaceFolders?.[0]?.uri?.fsPath || '';
|
|
xmlParts.push(workspaceRoot);
|
|
const fileTree = this.generateFileTree(files, workspaceRoot);
|
|
xmlParts.push(fileTree);
|
|
xmlParts.push('</file_map>');
|
|
}
|
|
|
|
xmlParts.push('<file_contents>');
|
|
|
|
for (const [filePath, { content, tokens }] of files) {
|
|
const extension = path.extname(filePath);
|
|
let language = extension.substring(1);
|
|
|
|
if (extension === '.js' || extension === '.jsx') {
|
|
language = 'javascript';
|
|
} else if (extension === '.ts' || extension === '.tsx') {
|
|
language = 'typescript';
|
|
} else if (extension === '.md') {
|
|
language = 'md';
|
|
} else if (extension === '.py') {
|
|
language = 'python';
|
|
} else if (extension === '.html') {
|
|
language = 'html';
|
|
} else if (extension === '.css') {
|
|
language = 'css';
|
|
} else if (extension === '.json') {
|
|
language = 'json';
|
|
}
|
|
|
|
const formattedContent = content;
|
|
const workspaceRoot = vscode.workspace.workspaceFolders?.[0]?.uri?.fsPath || '';
|
|
|
|
xmlParts.push(`File: ${this.getRelativePath(filePath, workspaceRoot)}`);
|
|
xmlParts.push(`\`\`\`${language}`);
|
|
xmlParts.push(formattedContent);
|
|
xmlParts.push('\`\`\`');
|
|
}
|
|
|
|
xmlParts.push('</file_contents>');
|
|
|
|
if (settings.tokenCalculationEnabled) {
|
|
const totalTokens = Array.from(files.values()).reduce((sum, { tokens }) => sum + tokens, 0);
|
|
vscode.window.showInformationMessage(`Total tokens: ${totalTokens}`);
|
|
}
|
|
|
|
if (settings.includeFormattingInstructions && formattingInstructions) {
|
|
xmlParts.push(formattingInstructions);
|
|
}
|
|
|
|
return xmlParts.join('\n');
|
|
}
|
|
|
|
/**
|
|
* Generate a tree representation of files
|
|
* @param files Map of file paths
|
|
* @param rootPath The workspace root path
|
|
* @returns String representation of the file tree
|
|
*/
|
|
private static generateFileTree(files: Map<string, any>, rootPath: string): string {
|
|
const treeLines: string[] = [];
|
|
|
|
interface TreeNode {
|
|
name: string;
|
|
isDirectory: boolean;
|
|
children: Map<string, TreeNode>;
|
|
}
|
|
|
|
const root: TreeNode = {
|
|
name: path.basename(rootPath),
|
|
isDirectory: true,
|
|
children: new Map<string, TreeNode>()
|
|
};
|
|
|
|
const ig = ignoreFunc();
|
|
try {
|
|
const gitignorePath = path.join(rootPath, '.gitignore');
|
|
if (fs.existsSync(gitignorePath)) {
|
|
const gitignoreContent = fs.readFileSync(gitignorePath, 'utf8');
|
|
ig.add(gitignoreContent);
|
|
ig.add('!.gitignore');
|
|
}
|
|
} catch (error) {
|
|
console.error('Error reading .gitignore:', error);
|
|
}
|
|
|
|
for (const filePath of files.keys()) {
|
|
const relativePath = this.getRelativePath(filePath, rootPath);
|
|
const normalizedPath = relativePath.split(path.sep).join('/');
|
|
|
|
if (ig.ignores(normalizedPath)) {
|
|
continue;
|
|
}
|
|
|
|
const parts = relativePath.split('/');
|
|
let currentNode = root;
|
|
|
|
for (let i = 0; i < parts.length; i++) {
|
|
const part = parts[i];
|
|
const isDirectory = i < parts.length - 1;
|
|
|
|
if (!currentNode.children.has(part)) {
|
|
currentNode.children.set(part, {
|
|
name: part,
|
|
isDirectory,
|
|
children: new Map<string, TreeNode>()
|
|
});
|
|
} else if (isDirectory) {
|
|
currentNode.children.get(part)!.isDirectory = true;
|
|
}
|
|
|
|
currentNode = currentNode.children.get(part)!;
|
|
}
|
|
}
|
|
|
|
const buildTreeLines = (node: TreeNode, prefix: string = '', isLast: boolean = true, parentPrefix: string = ''): void => {
|
|
if (node !== root) {
|
|
const linePrefix = parentPrefix + (isLast ? '└── ' : '├── ');
|
|
treeLines.push(`${linePrefix}${node.name}${node.isDirectory ? '' : ''}`);
|
|
}
|
|
|
|
const sortedChildren = Array.from(node.children.values())
|
|
.sort((a, b) => {
|
|
if (a.isDirectory === b.isDirectory) {
|
|
return a.name.localeCompare(b.name);
|
|
}
|
|
return a.isDirectory ? -1 : 1;
|
|
});
|
|
|
|
sortedChildren.forEach((child, index) => {
|
|
const isChildLast = index === sortedChildren.length - 1;
|
|
const childParentPrefix = node === root ? '' : parentPrefix + (isLast ? ' ' : '│ ');
|
|
buildTreeLines(child, prefix, isChildLast, childParentPrefix);
|
|
});
|
|
};
|
|
|
|
buildTreeLines(root);
|
|
return treeLines.join('\n');
|
|
}
|
|
|
|
/**
|
|
* Get the path relative to the workspace root
|
|
* @param filePath Absolute file path
|
|
* @param rootPath Workspace root path
|
|
* @returns Relative path
|
|
*/
|
|
private static getRelativePath(filePath: string, rootPath: string): string {
|
|
if (filePath.startsWith(rootPath)) {
|
|
const relativePath = filePath.substring(rootPath.length);
|
|
return relativePath.startsWith('/') ? relativePath.substring(1) : relativePath;
|
|
}
|
|
return filePath;
|
|
}
|
|
}
|