fix: various bugs

Author: rrrodzilla

src/docLoader.ts

@@ -1,5 +1,5 @@
 import { HTMLReader } from "@llamaindex/readers/html";
-import { glob } from "glob";
+import { glob, GlobOptions } from "glob";
 import path from "path";
 import fs from "fs/promises";
 import { Document } from "llamaindex";
@@ -15,51 +15,60 @@ export async function loadCrateDocs(
 ): Promise<Document<Metadata>[]> {
   const docsPath = path.join(rootDocsPath, crateName);
   const pattern: string = "**/*.html";
-  const options: import("glob").GlobOptions = {
+  const options: GlobOptions = {
     cwd: docsPath,
     ignore: "**/index.html",
     withFileTypes: false,
   };
 
-  const htmlFiles: string[] = await glob(pattern, options) as string[];
+  const htmlFiles: string[] = (await glob(pattern, options)) as string[];
   console.log("found # html files (initial):", htmlFiles.length);
 
-  // Group files by basename and count occurrences
-  const basenameGroups: { [key: string]: string[] } = {};
-  htmlFiles.forEach((file: string) => {
-    const basename: string = path.basename(file);
-    if (!basenameGroups[basename]) {
-      basenameGroups[basename] = [];
-    }
-    basenameGroups[basename].push(file);
-  });
+  // Check the dangerous flag
+  const includeAllDangerously = process.env.INCLUDE_ALL_DOCS_DANGEROUSLY === "true";
+  let filesToProcess: string[] = [];
+
+  if (includeAllDangerously) {
+    console.warn(
+      `WARNING: INCLUDE_ALL_DOCS_DANGEROUSLY is enabled for ${crateName}. Processing all ${htmlFiles.length} HTML files. This may incur significant time or cost (e.g., embedding calls).`
+    );
+    filesToProcess = htmlFiles;
+  } else {
+    // Group files by basename and count occurrences
+    const basenameGroups: { [key: string]: string[] } = {};
+    htmlFiles.forEach((file: string) => {
+      const basename: string = path.basename(file);
+      if (!basenameGroups[basename]) {
+        basenameGroups[basename] = [];
+      }
+      basenameGroups[basename].push(file);
+    });
 
-  // Filter for duplicate basenames and pick largest file
-  const uniqueDuplicateFiles: string[] = [];
-  for (const [basename, files] of Object.entries(basenameGroups)) {
-    if (files.length > 1) {
-      const largestFile = await files.reduce(
-        async (largestPromise, current) => {
-          const largest = await largestPromise;
-          const largestStats = await fs.stat(path.join(docsPath, largest));
-          const currentStats = await fs.stat(path.join(docsPath, current));
-          return largestStats.size > currentStats.size ? largest : current;
-        },
-        Promise.resolve(files[0])
-      );
-      uniqueDuplicateFiles.push(largestFile);
+    // Filter for duplicate basenames and pick largest file
+    for (const [basename, files] of Object.entries(basenameGroups)) {
+      if (files.length > 1) {
+        const largestFile = await files.reduce(
+          async (largestPromise, current) => {
+            const largest = await largestPromise;
+            const largestStats = await fs.stat(path.join(docsPath, largest));
+            const currentStats = await fs.stat(path.join(docsPath, current));
+            return largestStats.size > currentStats.size ? largest : current;
+          },
+          Promise.resolve(files[0])
+        );
+        filesToProcess.push(largestFile);
+      }
     }
+    console.log(
+      "found # unique duplicate html files for crate " + crateName + ":",
+      filesToProcess.length
+    );
   }
 
-  console.log(
-    "found # unique duplicate html files for crate " + crateName + ":",
-    uniqueDuplicateFiles.length
-  );
-
-  // Process the unique duplicate files with HTMLReader
+  // Process the selected files with HTMLReader
   const htmlReader: HTMLReader = new HTMLReader();
   const docs: Document<Metadata>[][] = await Promise.all(
-    uniqueDuplicateFiles.map(async (filePath: string) => {
+    filesToProcess.map(async (filePath: string) => {
       const fullFilePath: string = path.join(docsPath, filePath);
       const fileDocs: Document<Metadata>[] = await htmlReader.loadData(fullFilePath);
       console.log("Loaded:", filePath);
@@ -73,7 +82,7 @@ export async function loadCrateDocs(
   );
 
   const foundDocs: Document<Metadata>[] = docs.flat();
-  console.log("found # docs from unique duplicates:", foundDocs.length);
+  console.log("found # docs from processed files:", foundDocs.length);
   return foundDocs;
 }
 

src/index.ts

@@ -26,6 +26,7 @@ class RustDocsServer {
   private persistDir: string;
   private docsPath: string;
   private crateName: string;
+  private toolName: string;
 
   constructor(docsPath: string, crateName: string) {
     this.server = new Server(
@@ -44,24 +45,39 @@ class RustDocsServer {
     this.persistDir = path.join(__dirname, "storage", this.crateName);
     this.docsPath = docsPath;
     this.setupToolHandlers();
+    this.toolName = `query_rust_docs_${this.crateName}`;
   }
 
   private async initializeIndex() {
     try {
-      const indexExists = fs.existsSync(
-        path.join(this.persistDir, "vector_store.json")
-      );
-
-      const storageContext = await storageContextFromDefaults({
-        persistDir: this.persistDir,
-      });
+      const indexExists =
+        fs.existsSync(path.join(this.persistDir, "vector_store.json")) &&
+        fs.existsSync(path.join(this.persistDir, "doc_store.json")) &&
+        fs.existsSync(path.join(this.persistDir, "index_store.json"));
 
       if (indexExists) {
         console.log("Loading existing vector store from", this.persistDir);
+        const storageContext = await storageContextFromDefaults({
+          persistDir: this.persistDir,
+        });
         this.index = await VectorStoreIndex.init({ storageContext });
         return;
+      } else {
+        // Clear only the files inside the persist directory, not the directory itself
+        if (fs.existsSync(this.persistDir)) {
+          const files = fs.readdirSync(this.persistDir);
+          for (const file of files) {
+            fs.unlinkSync(path.join(this.persistDir, file));
+          }
+        } else {
+          // Ensure the directory exists if it was never created
+          fs.mkdirSync(this.persistDir, { recursive: true });
+        }
       }
 
+      const storageContext = await storageContextFromDefaults({
+        persistDir: this.persistDir,
+      });
       this.index = await VectorStoreIndex.fromDocuments(
         await loadCrateDocs(this.docsPath, this.crateName),
         {
@@ -80,7 +96,7 @@ class RustDocsServer {
     this.server.setRequestHandler(ListToolsRequestSchema, async () => ({
       tools: [
         {
-          name: "query_rust_docs",
+          name: this.toolName,
           description: `Query the official Rust documentation for the '${this.crateName}' crate. Use this tool to retrieve detailed information about '${this.crateName}'’s API, including structs, traits, enums, constants, and functions. Ideal for answering technical questions about how to use '${this.crateName}' in Rust projects, such as understanding specific methods, configuration options, or integration details. Additionally, leverage this tool to ensure accuracy of written code by verifying API usage and to resolve Clippy or lint errors by clarifying correct implementations. For example, use it for questions like "How do I configure routing in ${this.crateName}?", "What does this ${this.crateName} struct do?", "Is this ${this.crateName} method call correct?", or "How do I fix a Clippy warning about ${this.crateName} in my code?"`,
           inputSchema: {
             type: "object",
@@ -100,36 +116,36 @@ class RustDocsServer {
         },
       ],
     }));
-  
+
     this.server.setRequestHandler(
       CallToolRequestSchema,
       async (request: CallToolRequest) => {
-        if (request.params.name !== "query_rust_docs") {
+        if (request.params.name !== this.toolName) {
           throw new McpError(
             ErrorCode.MethodNotFound,
             `Unknown tool: ${request.params.name}`
           );
         }
-  
+
         const { question, crate } = request.params.arguments as {
           question: string;
           crate: string;
         };
-  
+
         if (crate !== this.crateName) {
           throw new McpError(
             ErrorCode.InvalidParams,
             `This server only supports queries for '${this.crateName}', not '${crate}'`
           );
         }
-  
+
         try {
           const queryEngine = this.index.asQueryEngine();
           const response: EngineResponse = await queryEngine.query({
             query: question,
             stream: false,
           });
-  
+
           return {
             content: [
               {
@@ -162,9 +178,12 @@ if (!process.env.CRATE_NAME) {
 if (!process.env.DOCS_PATH) {
   throw new Error("DOCS_PATH environment variable must be set");
 }
-if (!process.env.OPENAI_API_KEY) { 
+if (!process.env.OPENAI_API_KEY) {
   throw new Error("OPENAI_API_KEY environment variable must be set");
 }
 
-const server = new RustDocsServer(process.env.DOCS_PATH, process.env.CRATE_NAME);
+const server = new RustDocsServer(
+  process.env.DOCS_PATH,
+  process.env.CRATE_NAME
+);
 server.run().catch(console.error);