colthreepv · September 13, 2024 21:37
diff --git a/create-context.js b/create-context.js
 #!/usr/bin/env node

 const { readdirSync, statSync, readFileSync, writeFileSync, readSync, openSync, closeSync } = require('node:fs');
 const { join, extname, relative } = require('node:path');
 const { execSync } = require('node:child_process');
 const { argv, exit } = require('node:process');

 const DEFAULT_IGNORE_PATHS = new Set([
  '.git', '.gitignore', '.dockerignore', 'node_modules', 'dist', 'build', '.next',
  '.env', '.env.local', '.env.production.local', '.env.development.local',
  '.env.test.local', 'package-lock.json', 'pnpm-lock.yaml', '.editorconfig',
  'LICENSE', 'yarn.lock', 'pnpm-workspace.yaml'
 ]);

 function isJSONFile(filePath) {
  return extname(filePath).toLowerCase() === '.json';
 }

 function isBinaryFile(filePath) {
  const fd = openSync(filePath, 'r');
  const buffer = Buffer.alloc(4096); // Read first 4KB

  try {
    const bytesRead = readSync(fd, buffer, 0, 4096, 0);
    closeSync(fd);

    // Check for NULL bytes in the first 4KB
    for (let i = 0; i < bytesRead; i++) {
      if (buffer[i] === 0) return true;
    }

    return false;
  } catch (error) {
    console.error(`Error reading file ${filePath}:`, error.message);
    return false;
  }
 }

 function truncateJSON(content) {
  try {
    const parsed = JSON.parse(content);
    if (Array.isArray(parsed)) {
      return JSON.stringify(parsed.slice(0, 2), null, 2);
    } else if (typeof parsed === 'object' && parsed !== null) {
      const truncated = Object.fromEntries(
        Object.entries(parsed).slice(0, 4)
      );
      return JSON.stringify(truncated, null, 2);
    }
    return content; // Return original content if it's not an array or object
  } catch (error) {
    console.error('Error parsing JSON:', error.message);
    return content; // Return original content if parsing fails
  }
 }

 function readFilesInDirectory(directoryPath, basePath, additionalPaths = []) {
  const ignorePaths = new Set([...DEFAULT_IGNORE_PATHS, ...additionalPaths]);
  let context = '';

  const files = readdirSync(directoryPath);
  for (const file of files) {
    if (ignorePaths.has(file)) continue;
    const filePath = join(directoryPath, file);
    const stats = statSync(filePath);
    const relativePath = relative(basePath, filePath);

    if (stats.isDirectory()) {
      context += readFilesInDirectory(filePath, basePath, additionalPaths);
      continue;
    }

    // conditions to exclude files
    if (file.startsWith('llm-context')) continue;
    if (isBinaryFile(filePath)) continue;

    // base case
    let content = readFileSync(filePath, 'utf8');
    if (isJSONFile(filePath)) content = truncateJSON(content);
    context += `\n# ${relativePath}\n\n${content}\n`;
  }

  return context;
 }

 // Function to get tree command output
 function getTreeOutput(directoryPath, additionalPaths = []) {
  const ignorePaths = Array.from(new Set([...DEFAULT_IGNORE_PATHS, ...additionalPaths]));
  try {
    // Execute tree command, ignoring specified paths
    const treeOutput = execSync(`tree -I '${ignorePaths.join('|')}' ${directoryPath}`, { encoding: 'utf8' });
    return `\n# Directory Structure\n\n${treeOutput}\n`;
  } catch (error) {
    console.error('Error executing tree command:', error.message);
    return '\n# Directory Structure\n\nUnable to generate directory structure.\n';
  }
 }

 // Main function
 function createContextFile(directoryPath, outputPath) {
  const args = argv.slice(4);
  if (args.length > 0) console.log('Additional paths to ignore:', args);
  let context = getTreeOutput(directoryPath, args);
  context += readFilesInDirectory(directoryPath, directoryPath, args);

  // Write the context to the output file
  writeFileSync(outputPath, context);
  console.log(`Context file created at: ${outputPath}`);
 }

 // Usage
 const providedDirPath = argv[2];
 const outputSuffix = argv[3];
 if (providedDirPath == null || outputSuffix == null) {
  console.log(`
 USAGE:
  node create-context.js <directoryPath> <outputSuffix> [ignorePath0] [ignorePath1] [ignorePath2]
 `);
  exit(1);
 }

 const directoryPath = join(process.cwd(), providedDirPath);
 const outputPath = `llm-context.${outputSuffix}.txt`;

 createContextFile(directoryPath, outputPath);