docs-loader.ts

import { logger, UUID } from '@elizaos/core';
import * as fs from 'fs';
import * as path from 'path';
import { KnowledgeService } from './service.ts';
import { AddKnowledgeOptions } from './types.ts';
import { isBinaryContentType } from './utils.ts';

/**
 * Get the knowledge path from runtime settings, environment, or default to ./docs
 */
export function getKnowledgePath(runtimePath?: string): string {
  // Priority: runtime setting > environment variable > default
  const knowledgePath =
    runtimePath || process.env.KNOWLEDGE_PATH || path.join(process.cwd(), 'docs');
  const resolvedPath = path.resolve(knowledgePath);

  if (!fs.existsSync(resolvedPath)) {
    logger.warn({ src: 'plugin:knowledge', knowledgePath: resolvedPath }, 'Knowledge path does not exist');
    if (runtimePath) {
      logger.warn({ src: 'plugin:knowledge' }, 'Please create the directory or update KNOWLEDGE_PATH in agent settings');
    } else if (process.env.KNOWLEDGE_PATH) {
      logger.warn({ src: 'plugin:knowledge' }, 'Please create the directory or update KNOWLEDGE_PATH environment variable');
    } else {
      logger.info({ src: 'plugin:knowledge' }, 'To use the knowledge plugin, create a docs folder or set KNOWLEDGE_PATH');
    }
  }

  return resolvedPath;
}

/**
 * Load documents from the knowledge path
 */
export async function loadDocsFromPath(
  service: KnowledgeService,
  agentId: UUID,
  worldId?: UUID,
  knowledgePath?: string
): Promise<{ total: number; successful: number; failed: number }> {
  const docsPath = getKnowledgePath(knowledgePath);

  if (!fs.existsSync(docsPath)) {
    logger.warn({ src: 'plugin:knowledge', agentId: agentId, docsPath }, 'Knowledge path does not exist');
    return { total: 0, successful: 0, failed: 0 };
  }

  logger.info({ src: 'plugin:knowledge', agentId: agentId, docsPath }, 'Loading documents from path');

  // Get all files recursively
  const files = getAllFiles(docsPath);

  if (files.length === 0) {
    logger.info({ src: 'plugin:knowledge', agentId: agentId }, 'No files found in knowledge path');
    return { total: 0, successful: 0, failed: 0 };
  }

  logger.info({ src: 'plugin:knowledge', agentId: agentId, filesCount: files.length }, 'Found files to process');

  let successful = 0;
  let failed = 0;

  for (const filePath of files) {
    try {
      const fileName = path.basename(filePath);
      const fileExt = path.extname(filePath).toLowerCase();

      // Skip hidden files and directories
      if (fileName.startsWith('.')) {
        continue;
      }

      // Determine content type
      const contentType = getContentType(fileExt);

      // Skip unsupported file types
      if (!contentType) {
        logger.debug({ src: 'plugin:knowledge', agentId: agentId, filePath }, 'Skipping unsupported file type');
        continue;
      }

      // Read file
      const fileBuffer = fs.readFileSync(filePath);

      // Check if file is binary using the same logic as the service
      const isBinary = isBinaryContentType(contentType, fileName);

      // For text files, read as UTF-8 string directly
      // For binary files, convert to base64
      const content = isBinary ? fileBuffer.toString('base64') : fileBuffer.toString('utf-8');

      // Create knowledge options
      const knowledgeOptions: AddKnowledgeOptions = {
        clientDocumentId: '' as UUID, // Will be generated by the service based on content
        contentType,
        originalFilename: fileName,
        worldId: worldId || agentId,
        content,
        roomId: agentId,
        entityId: agentId,
      };

      // Process the document
      logger.debug({ src: 'plugin:knowledge', agentId: agentId, filename: fileName }, 'Processing document');
      const result = await service.addKnowledge(knowledgeOptions);

      logger.info({ src: 'plugin:knowledge', agentId: agentId, filename: fileName, fragmentCount: result.fragmentCount }, 'Document processed');
      successful++;
    } catch (error: any) {
      logger.error({ src: 'plugin:knowledge', agentId: agentId, filePath, error: error.message }, 'Failed to process file');
      failed++;
    }
  }

  logger.info(
    { src: 'plugin:knowledge', agentId: agentId, successful, failed, total: files.length },
    'Document loading complete'
  );

  return {
    total: files.length,
    successful,
    failed,
  };
}

/**
 * Recursively get all files in a directory
 */
function getAllFiles(dirPath: string, files: string[] = []): string[] {
  try {
    const entries = fs.readdirSync(dirPath, { withFileTypes: true });

    for (const entry of entries) {
      const fullPath = path.join(dirPath, entry.name);

      if (entry.isDirectory()) {
        // Skip node_modules and other common directories
        if (!['node_modules', '.git', '.vscode', 'dist', 'build'].includes(entry.name)) {
          getAllFiles(fullPath, files);
        }
      } else if (entry.isFile()) {
        files.push(fullPath);
      }
    }
  } catch (error: any) {
    logger.error({ src: 'plugin:knowledge', dirPath, error: error.message }, 'Error reading directory');
  }

  return files;
}

/**
 * Get content type based on file extension
 */
function getContentType(extension: string): string | null {
  const contentTypes: Record<string, string> = {
    // Text documents
    '.txt': 'text/plain',
    '.md': 'text/markdown',
    '.markdown': 'text/markdown',
    '.tson': 'text/plain',
    '.xml': 'application/xml',
    '.csv': 'text/csv',
    '.tsv': 'text/tab-separated-values',
    '.log': 'text/plain',

    // Web files
    '.html': 'text/html',
    '.htm': 'text/html',
    '.css': 'text/css',
    '.scss': 'text/x-scss',
    '.sass': 'text/x-sass',
    '.less': 'text/x-less',

    // JavaScript/TypeScript
    '.js': 'text/javascript',
    '.jsx': 'text/javascript',
    '.ts': 'text/typescript',
    '.tsx': 'text/typescript',
    '.mjs': 'text/javascript',
    '.cjs': 'text/javascript',
    '.vue': 'text/x-vue',
    '.svelte': 'text/x-svelte',
    '.astro': 'text/x-astro',

    // Python
    '.py': 'text/x-python',
    '.pyw': 'text/x-python',
    '.pyi': 'text/x-python',

    // Java/Kotlin/Scala
    '.java': 'text/x-java',
    '.kt': 'text/x-kotlin',
    '.kts': 'text/x-kotlin',
    '.scala': 'text/x-scala',

    // C/C++/C#
    '.c': 'text/x-c',
    '.cpp': 'text/x-c++',
    '.cc': 'text/x-c++',
    '.cxx': 'text/x-c++',
    '.h': 'text/x-c',
    '.hpp': 'text/x-c++',
    '.cs': 'text/x-csharp',

    // Other languages
    '.php': 'text/x-php',
    '.rb': 'text/x-ruby',
    '.go': 'text/x-go',
    '.rs': 'text/x-rust',
    '.swift': 'text/x-swift',
    '.r': 'text/x-r',
    '.R': 'text/x-r',
    '.m': 'text/x-objectivec',
    '.mm': 'text/x-objectivec',
    '.clj': 'text/x-clojure',
    '.cljs': 'text/x-clojure',
    '.ex': 'text/x-elixir',
    '.exs': 'text/x-elixir',
    '.lua': 'text/x-lua',
    '.pl': 'text/x-perl',
    '.pm': 'text/x-perl',
    '.dart': 'text/x-dart',
    '.hs': 'text/x-haskell',
    '.elm': 'text/x-elm',
    '.ml': 'text/x-ocaml',
    '.fs': 'text/x-fsharp',
    '.fsx': 'text/x-fsharp',
    '.vb': 'text/x-vb',
    '.pas': 'text/x-pascal',
    '.d': 'text/x-d',
    '.nim': 'text/x-nim',
    '.zig': 'text/x-zig',
    '.jl': 'text/x-julia',
    '.tcl': 'text/x-tcl',
    '.awk': 'text/x-awk',
    '.sed': 'text/x-sed',

    // Shell scripts
    '.sh': 'text/x-sh',
    '.bash': 'text/x-sh',
    '.zsh': 'text/x-sh',
    '.fish': 'text/x-fish',
    '.ps1': 'text/x-powershell',
    '.bat': 'text/x-batch',
    '.cmd': 'text/x-batch',

    // Config files
    '.json': 'application/json',
    '.yaml': 'text/x-yaml',
    '.yml': 'text/x-yaml',
    '.toml': 'text/x-toml',
    '.ini': 'text/x-ini',
    '.cfg': 'text/x-ini',
    '.conf': 'text/x-ini',
    '.env': 'text/plain',
    '.gitignore': 'text/plain',
    '.dockerignore': 'text/plain',
    '.editorconfig': 'text/plain',
    '.properties': 'text/x-properties',

    // Database
    '.sql': 'text/x-sql',

    // Binary documents
    '.pdf': 'application/pdf',
    '.doc': 'application/msword',
    '.docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
  };

  return contentTypes[extension] || null;
}