docs-agent 1.1.0

package/src/DocsAgent.js

@@ -0,0 +1,728 @@
+import ReviewPrompt from './config/prompt.review.js';
+import PrioritizePrompt from './config/prompt.prioritize.js';
+import EditPrompt from './config/prompt.edit.js';
+import DisruptiveEditPrompt from './config/prompt.edit.disruptive.js';
+import LinkingPrompt from './config/prompt.linking.js';
+import GenerateReferenceDocsPrompt from './config/prompt.gen.referencedocs.js';
+import DocsVsCodePrompt from './config/prompt.docs.vs.code.js';
+import ExtractCodeAndSymbolsPrompt from './config/prompt.extract.code.js';
+import LLM from './LLM.js';
+import * as FileUtility from './FileUtility.js';
+import CodeSearch from './CodeSearch.js';
+import { z } from "zod";
+
+class DocsAgent {
+  constructor(accessMode = 'mcp') {
+    this.name = 'Docs Agent';
+    this.accessMode = accessMode;
+  }
+
+  /**
+   * Plan and execute the docs improvement process for the given docs page content
+   * @param {string} content - The content of the docs page to improve
+   * @param {object} context - The context of the docs page
+   * @param {string} context.filepath - The absolute filepath of the docs page to improve
+   * @param {string} context.filename - The filename of the docs page
+   * @param {string} context.projectStructure - The project structure as folders/files hierarchy
+   * @param {string} customInstructions - Custom instructions for the docs improvement process
+   * @param {string} context.glossaryFile - The filepath of the glossary file, containing a list of project-specific terms and their definitions, used for consistent and technically accurate documentation
+   * @param {string[]} context.relatedFiles - The filepaths of the docs pages that are related to the current page being reviewed/edited
+   * @param {boolean} isFileWriteAllowed - Whether to write the edited content to the file system
+   * @returns {Promise<string>} - The edited docs content
+   */
+  async reviewPrioritizeAndEdit(content, context={}, customInstructions, isFileWriteAllowed){
+    console.log("Starting the docs improvement process now");
+    if(!content){
+      if(!context?.filepath){
+        throw new Error("Content or filepath is required");
+      }
+      content = await FileUtility.readFile(context.filepath, this.accessMode);
+    }
+    console.log("Received content to review, prioritize and edit");
+    console.log(context?.filename, context?.projectStructure, context?.glossaryFile, context?.relatedFiles);
+    if(!context?.filename){
+      context.filename = "docs.md";
+    }
+    if(context?.glossaryFile){
+      context.glossary = await FileUtility.readFile(context.glossaryFile, this.accessMode);
+    }
+    if(context?.relatedFiles){
+      // TODO: Remove duplicates
+      context.relatedFilesContent = await Promise.all(context.relatedFiles.map(async (file) => {
+        return FileUtility.readFile(file, this.accessMode);
+      }));
+    }
+    // Save the original content
+    if(isFileWriteAllowed){
+      await FileUtility.writeFile(path.join(process.cwd(), "public", "original", context.filename), content, this.accessMode);
+    }
+    // Review the content
+    const review = await this.review(content, context, customInstructions);
+
+    // Prioritize the editing tasks
+    const prioritize = await this.prioritize(content, review, context, customInstructions);
+
+    // Edit the content as per the prioritized tasks
+    const editedContent = await this.edit(content, prioritize, context, customInstructions);
+
+    // Save the edited content
+    if(context.filename && isFileWriteAllowed){
+      await FileUtility.writeFile(path.join(process.cwd(), "public", context.filename), editedContent, this.accessMode);
+    }
+    return editedContent;
+  }
+
+  /**
+   * Review the documentation and suggest improvements.
+   * @param {string} content - The documentation to review.
+   * @param {string} context - The context of the documentation.
+   * @param {string} customInstructions - Custom instructions for the docs improvement process
+   * @returns {Promise<string>} - The review of the documentation.
+   */
+  async review(content, context, customInstructions) {
+    console.log("Reviewing the docs", content?.slice(0, 300)+"...");
+    if(!content && context?.filepath){
+      content = await FileUtility.readFile(context.filepath, this.accessMode);
+    }
+    if(!content){
+      throw new Error("Content is required");
+    }
+    const llm = new LLM({
+      aiService: process.env.REVIEW_AI_SERVICE || "gemini",
+      model: process.env.REVIEW_AI_MODEL || "gemini-2.5-pro-exp-03-25"
+    });
+    let userMessage = "Content:\n"+content;
+    if(context?.projectStructure){
+      userMessage = userMessage + "\n\n\n\n\n\nProject Structure:\n"+context.projectStructure;
+    }
+    if(context?.filepath){
+      userMessage = userMessage + "\n\n\n\n\n\nPath of the file being reviewed (the content is already provided): \n"+context.filepath + " [Ensure to mention the relative path (from the root of the github repository project only) to this file in the output (the review result)]";
+    }
+    const review = await llm.chat([
+      {
+        role: "system",
+        content: ReviewPrompt + "\n\n\n\n\n\nCustom Instructions:\n" + customInstructions
+      },
+      {
+        role: "user",
+        content: [
+          {
+          type: "text",
+          text: userMessage
+        }
+      ]
+    }], null, {
+      isEnabled: true,
+      functionId: "reviewDocs",
+      metadata: {
+        customInstructions: customInstructions,
+      }
+    });
+    console.log("Review:\n"+review?.slice(0, 600)+"...");
+    return review;
+  }
+
+  /**
+   * Break down review into a documentation editing tasks and prioritize them based on their impact vs disruption..
+   * @param {string} review - The review of the documentation.
+   * @param {string} content - The documentation to review.
+   * @param {string} context - The context of the documentation.
+   * @param {string} customInstructions - Custom instructions for the docs improvement process
+   * @returns {Promise<string>} - The prioritized documentation editing plan.
+   */
+  async prioritize(content, review, context = {}, customInstructions) {
+    console.log("Prioritizing the docs improvement plan", review);
+    if(!content && context?.filepath){
+      content = await FileUtility.readFile(context.filepath, this.accessMode);
+    }
+    if(!content || !review){
+      throw new Error("Content and review are required");
+    }
+    const llm = new LLM({
+      aiService: "anthropic",
+      model: "claude-3-5-sonnet-20240620",
+      temperature: 0,
+      topP: 0.95
+    });
+    // System prompt
+    const ALLOW_DISRUPTIVE_CHANGES = (process.env.ALLOW_DISRUPTIVE_CHANGES === "true") || context?.allowDisruptiveChanges;
+    let systemPrompt = PrioritizePrompt;
+    systemPrompt += `\n\nRestructuring the project files/folder structure is ${ALLOW_DISRUPTIVE_CHANGES ? "allowed" : "NOT ALLOWED"}.`;
+    if(customInstructions){
+      systemPrompt += "\n\n\n\n\n\nCustom Instructions:\n" + customInstructions;
+    }
+    // User message
+    let userMessage = "Content:\n"+content;
+    if(context?.projectStructure){
+      userMessage = userMessage + "\n\n\n\n\n\nProject Structure:\n"+context.projectStructure;
+    }
+    if(review){
+      userMessage = userMessage + "\n\n\n\n\n\nContent Review By The Documentation Expert:\n"+review;
+    }
+    const prioritized = await llm.chat([
+      {
+        role: "system",
+        content: systemPrompt
+      },
+      {
+        role: "user",
+        content: [
+          {
+            type: "text",
+            text: userMessage
+          }
+        ]
+      }
+    ], null, {
+      isEnabled: true,
+      functionId: "prioritizeDocsEditTasks",
+      metadata: {
+        customInstructions: customInstructions,
+      }
+    });
+    console.log(prioritized);
+    return prioritized;
+  }
+
+  /**
+   * Edit the documentation.
+   * @param {string} content - The documentation to edit.
+   * @param {string} editPlan - The edit plan. Specific changes to be made to the documentation.
+   * @param {string} context - The context of the documentation.
+   * @param {string} context.filepath - The absolute filepath of the docs page to edit
+   * @param {string} context.filename - The filename of the docs page
+   * @param {string} context.projectStructure - The project structure as folders/files hierarchy
+   * @param {string} context.glossaryFile - The filepath of the glossary file, containing a list of project-specific terms and their definitions, used for consistent and technically accurate documentation
+   * @param {string[]} context.relatedFiles - The filepaths of the docs pages that are related to the current page being reviewed/edited
+   * @param {string} customInstructions - Custom instructions for the docs editing process
+   * @returns {Promise<string>} - The edited documentation.
+   */
+  async edit(content, editPlan, context = {}, customInstructions) {
+    console.log("Editing the docs according to the edit plan", editPlan);
+    if(!editPlan){
+      throw new Error("Edit plan is required");
+    }
+    if(!content && context?.filepath){
+      content = await FileUtility.readFile(context.filepath, this.accessMode);
+    }
+    if(!content){
+      throw new Error("Content is required");
+    }
+    const llm = new LLM({
+      aiService: "anthropic",
+      model: "claude-3-5-sonnet-20240620",
+      temperature: 0,
+      topP: 0.95,
+      maxInputTokens: 50000
+    });
+    // System prompt
+    const ALLOW_DISRUPTIVE_CHANGES = (process.env.ALLOW_DISRUPTIVE_CHANGES === "true") || context?.allowDisruptiveChanges;
+    let systemPrompt = ALLOW_DISRUPTIVE_CHANGES ? DisruptiveEditPrompt : EditPrompt;
+    if(customInstructions){
+      systemPrompt = systemPrompt + "\n\n\n\n\n\nCustom Instructions:\n" + customInstructions;
+    }
+    // User message
+    let userMessage = "Content:\n"+content;
+    if(context?.projectStructure){
+      userMessage = userMessage + "\n\n\n\n\n\nProject Structure:\n"+context.projectStructure;
+    }
+    if(editPlan){
+      userMessage = userMessage + "\n\n\n\n\n\nEdit Plan:\n"+editPlan;
+    }
+    // Assistant reassurance
+    let assistantReassurancePrompt = "I will edit the documentation strictly following the given principles. I will not add any other text or comments other than the main file content. I won't add the text 'Here's the edited version...' either. I will only return the main file content.";
+    // Optional context
+    let tokenEstimate = llm.estimateTokens(userMessage);
+    if(tokenEstimate < llm.maxInputTokens){
+      if(!context?.glossary && context?.glossaryFile){
+        try{
+          context.glossary = await FileUtility.readFile(context.glossaryFile, this.accessMode);
+        } catch(error){
+          console.error("Error reading glossary file", context.glossaryFile, error);
+        }
+      }
+      if(context?.glossary){
+        tokenEstimate += llm.estimateTokens(context.glossary);
+        if(tokenEstimate < llm.maxInputTokens){
+          assistantReassurancePrompt += "\n\nI will use following glossary to make the writing more consistent and professional:\n" + context.glossary;
+        }
+      }
+      if(!context?.relatedFilesContent && context?.relatedFiles){
+        context.relatedFilesContent = "";
+        for(const file of context.relatedFiles){
+          try{
+            const fileContent = await FileUtility.readFile(file, this.accessMode);
+            if(fileContent){
+              context.relatedFilesContent += "// "+file+"\n\n"+fileContent+"\n\n\n\n\n\n";
+            }
+          } catch(error){
+            console.error("Error reading related file", file, error);
+          }
+        }
+      }
+      if(context?.relatedFilesContent){
+        tokenEstimate += llm.estimateTokens(context.relatedFilesContent);
+        if(tokenEstimate < llm.maxInputTokens){
+          assistantReassurancePrompt += "\n\nI will use following related files content for proper context and linking:\n" + context.relatedFilesContent;
+        }
+      }
+    }
+    const edited = await llm.chat([
+      {
+        role: "system",
+        content: systemPrompt
+      },
+      {
+        role: "assistant",
+        content: assistantReassurancePrompt
+      },
+      {
+        role: "user",
+        content: [
+          {
+            type: "text",
+            text: userMessage
+          }
+        ]
+      }
+    ], null, {
+      isEnabled: true,
+      functionId: "editDocs",
+      metadata: {
+        customInstructions: customInstructions,
+        editPlan: editPlan,
+        allowDisruptiveChanges: ALLOW_DISRUPTIVE_CHANGES
+      }
+    });
+    return edited;
+  }
+
+  /**
+   * Improves internal linking of the docs page to the related files and glossary concepts.
+   * @param {string} content - The content of the docs page to linkify
+   * @param {object} context - The context of the docs page
+   * @param {string} context.filepath - The absolute filepath of the docs page to linkify
+   * @param {string} context.filename - The filename of the docs page
+   * @param {string} context.projectStructure - The project structure as folders/files hierarchy
+   * @param {string} customInstructions - Custom instructions for the docs improvement process
+   * @param {string} context.glossaryFile - The filepath of the glossary file, containing a list of project-specific terms and their definitions, used for consistent and technically accurate documentation   
+   * @param {string[]} context.relatedFiles - The filepaths of the docs pages that are related to the current page being reviewed/edited
+   * @param {string[]} context.referenceSourceCodeFiles - The remote URLs of the reference source code files for which the docs are being generated. Use this to verify and generate technically accurate documentation
+   * @returns {Promise<string>} - The edited docs content with improved internal linking
+   */
+  async linkify(content, context, customInstructions){
+    console.log("Linkifying the docs", content);
+    if(!content && context?.filepath){
+      content = await FileUtility.readFile(context.filepath, this.accessMode);
+    }
+    const llm = new LLM({
+      aiService: "anthropic",
+      model: "claude-3-5-sonnet-20240620",
+      temperature: 0,
+      topP: 0.95,
+      maxInputTokens: 80000
+    });
+    // System prompt
+    let systemPrompt = LinkingPrompt;
+    if(customInstructions){
+      systemPrompt = systemPrompt + "\n\n\n\n\n\nCustom Instructions:\n" + customInstructions;
+    }
+    // User message
+    let userMessage = "Content:\n"+content;
+    if(context?.projectStructure){
+      userMessage = userMessage + "\n\n\n\n\n\nProject Structure:\n"+context.projectStructure;
+    }
+    if(context?.relatedFilesContent){
+      if(context?.relatedFiles){
+        userMessage = userMessage + "\n\n\n\n\n\nRelated Files:\n"+context.relatedFiles.join("\n");
+      }
+      if(context?.relatedFilesContent){
+        userMessage = userMessage + "\n\n\n\n\n\nRelated Files Content:\n"+context.relatedFilesContent.join("\n\n");
+      }
+    }
+    if(context?.referenceSourceCodeFiles){
+      let referenceSourceCodeContent = "\n\n\n\n\n\nReference Source Code:\n";
+      for(const fileUrl of context.referenceSourceCodeFiles){
+        const fileContent = await FileUtility.readFile(fileUrl, this.accessMode);
+        userMessage += "// "+fileUrl+"\n\n"+fileContent+"\n\n\n\n\n\n";
+      }
+      userMessage += referenceSourceCodeContent;
+    }
+    if(context?.glossary){
+      if(context?.glossaryFile){
+        userMessage = userMessage + "\n\n\n\n\n\nGlossary File:\n"+context.glossaryFile;
+      }
+      if(context?.glossary){
+        userMessage = userMessage + "\n\n\n\n\n\nGlossary:\n"+context.glossary;
+      }
+    }
+    // Assistant reassurance
+    let assistantReassurancePrompt = "I will improve internal linking of the documentation strictly following the given principles. I will not add any other text or comments in the final output. I won't add the text 'Here's the edited content...' in the final output. I will only return the edited file content.";
+    const contentWithProperLinking = await llm.chat([
+      {
+        role: "system",
+        content: systemPrompt
+      },
+      {
+        role: "assistant",
+        content: assistantReassurancePrompt
+      },
+      {
+        role: "user",
+        content: [
+          {
+            type: "text",
+            text: userMessage
+          }
+        ]
+      }
+    ], null, {
+      isEnabled: true,
+      functionId: "linkifyDocs",
+      metadata: {
+        customInstructions: customInstructions,
+      }
+    });
+    return contentWithProperLinking;
+  }
+
+  /**
+   * Extract code and related symbols from docs
+   * @param {string} content - The content of the docs page to extract code and symbols from
+   * @param {object} context - The context of the docs page
+   * @param {string} customInstructions - Custom instructions to extract code and symbols
+   * @returns {Promise<{symbols: string[], codeBlocks: string[]}>} - The extracted code and related symbols
+   */
+  async extractCodeAndSymbols(content, context, customInstructions){
+    console.log("Extracting code and symbols from the docs", content);
+    const llm = new LLM({
+      aiService: "anthropic",
+      model: "claude-3-5-sonnet-20240620",
+      temperature: 0,
+      topP: 0.95,
+      maxInputTokens: 30000
+    });
+    let systemPrompt = ExtractCodeAndSymbolsPrompt;
+    if(customInstructions){
+      systemPrompt = systemPrompt + "\n\n\n\n\n\nCustom Instructions:\n" + customInstructions;
+    }
+    // User message
+    let userMessage = "Docs Page Content:\n"+content;
+    const { symbols, codeBlocks } = await llm.chat([
+      {
+        role: "system",
+        content: systemPrompt
+      },
+      {
+        role: "assistant",
+        content: "I will extract the code and related symbols from the docs page content. I will return the code and related symbols in a JSON format. I will not add any other text or comments in the final output. I won't add the text 'Here's the extracted code and symbols...' in the final output. I will only return the JSON object."
+      },
+      {
+        role: "user",
+        content: [
+          {
+            type: "text",
+            text: userMessage
+          }
+        ]
+      }
+    ], {
+      responseFormat: "json",
+      schema: z.object({
+        symbols: z.array(z.string()).describe("The symbols extracted from the docs page content"),
+        codeBlocks: z.array(z.string()).describe("The code blocks extracted from the docs page content")
+      })
+    }, {
+      isEnabled: true,
+      functionId: "extractCodeAndSymbols",
+      metadata: {
+        customInstructions: customInstructions,
+      }
+    });
+    console.log("Extracted "+symbols?.length+" symbols and "+codeBlocks?.length+" code blocks");
+    console.log("Symbols:", "\n    "+symbols?.join("\n    "));
+    console.log("Code Blocks:", "\n    "+codeBlocks?.join("\n    "));
+    return { symbols, codeBlocks };
+  }
+
+  /**
+   * Find the reference source code files for the given docs page content
+   * @param {string} content 
+   * @param {object} context 
+   * @param {string} context.repoUrl - GitHub repository URL
+   * @param {string} customInstructions - Custom instructions to extract symbols
+   * @returns {Promise<string>} - The edited docs content with improved internal linking
+   */
+  async findReferenceSourceCodeFiles(content, context, customInstructions){
+    // Extract symbols and code blocks from the docs page
+    let symbols = [];
+    let codeBlocks = [];
+    
+    try {
+      const extractionResult = await this.extractCodeAndSymbols(content, context, customInstructions);
+      symbols = extractionResult.symbols || [];
+      codeBlocks = extractionResult.codeBlocks || [];
+    } catch (error) {
+      console.error('Failed to extract code and symbols:', error);
+      console.log('Continuing with empty symbols and code blocks due to extraction failure');
+    }
+    
+    // Search for the symbols in the reference codebase
+    const codeSearch = new CodeSearch({
+      searchSpace: "github",
+      repository: context.repoUrl,
+      fetchContent: false,
+      maxResults: 5,
+      scoreThreshold: 0
+    });
+    // Sanitize search results
+    let referenceSourceCodeFilesMap = new Map();
+    const setOfSymbols = new Set(symbols);
+    let notFoundSymbols = new Set();
+    let foundSymbols = new Set();
+    for(const symbol of setOfSymbols){
+      const searchResults = await codeSearch.execute(symbol);
+      for(const result of searchResults){
+        let mapItemValue = referenceSourceCodeFilesMap.get(result.path) || {};
+        mapItemValue.results = result.results;
+        mapItemValue.symbols = [...(mapItemValue.symbols || []), symbol];
+        mapItemValue.totalScore = (mapItemValue.totalScore || 0) + result.score;
+        mapItemValue.path = result.path;
+        mapItemValue.filepath = result.filepath;
+        referenceSourceCodeFilesMap.set(result.path, mapItemValue);
+        foundSymbols.add(symbol);
+      }
+      if(!foundSymbols.has(symbol)){
+        notFoundSymbols.add(symbol);
+      }
+    }
+    // Sort reference files by aggregated score
+    const sortedReferenceSourceCodeFiles = Array.from(referenceSourceCodeFilesMap.values()).sort((a, b) => b.totalScore - a.totalScore);
+    console.log("Sorted "+sortedReferenceSourceCodeFiles?.length+" reference source code files:", "\n    "+sortedReferenceSourceCodeFiles?.map(result => (result.filepath || result.path) + " (Score: " + result.totalScore + ")").join("\n    "));
+    // Draft a summary
+    let summary = ""; // Initialize as empty string
+    if(notFoundSymbols.size > 0){
+      summary += "\n\n"+notFoundSymbols.size+" symbols not found in the codebase: "+Array.from(notFoundSymbols).join(", ");
+    }
+    if(foundSymbols.size > 0){
+      summary += "\n\n"+foundSymbols.size+" symbols found in the codebase: "+Array.from(foundSymbols).join(", ");
+      if(sortedReferenceSourceCodeFiles?.length > 0){
+        summary += "\n\nReference source code files where these symbols are found:\n    " +
+          sortedReferenceSourceCodeFiles?.map(result => {
+            const symbolsList = result.symbols ? result.symbols.join(", ") : "";
+            return (result.filepath || result.path) + " (Score: " + result.totalScore + "; Symbols: " + symbolsList + ")";
+          }).join("\n    ");
+      }
+    } else {
+      console.error("No symbols found in the codebase. Check if the repo url is correct and the repository is accessible.");
+    }
+    // TODO: Code blocks check (critical in case the docs does not have any symbols)
+    console.log("Reference File Search Summary:\n", summary);
+    return {
+      files: sortedReferenceSourceCodeFiles?.map(result => result.filepath || result.path),
+      fileScoreMap: referenceSourceCodeFilesMap,
+      summary
+    };
+  }
+
+  /**
+   * Verify if the docs are accurate and up to date with the source code
+   * @param {string} docsContentFilePath - The absolute filepath of the docs page to audit
+   * @param {object} context - The context of the docs page
+   * @param {string} context.repoUrl - GitHub repository URL
+   * @param {string} context.projectStructure - A nested markdown list of the project structure showing the folders/files hierarchy of the repository
+   * @param {string} customInstructions - Custom instructions for the audit process
+   * @returns {Promise<string>} - The audit report
+   */
+  async auditDocsAgainstCode(docsContentFilePath, context, customInstructions){
+    console.log("Comparing the docs with the source code to find the incorrect information", docsContentFilePath);
+    if(!docsContentFilePath){
+      throw new Error("docsContentFilePath is required");
+    }
+    // Docs content
+    const docsContent = await FileUtility.readFile(docsContentFilePath, this.accessMode);
+    if(!docsContent){ 
+      throw new Error("Couldn't read the docs content from the file: "+docsContentFilePath);
+    }
+    const { files: referenceSourceCodeFiles, fileScoreMap: referenceSourceCodeFilesMap, summary: referenceFilesSummary } = await this.findReferenceSourceCodeFiles(docsContent, context, customInstructions);
+    const llm = new LLM({
+      aiService: process.env.REVIEW_AI_SERVICE || "gemini",
+      model: process.env.REVIEW_AI_MODEL || "gemini-2.5-pro-preview-06-05"
+    });
+    // System prompt
+    let systemPrompt = DocsVsCodePrompt;
+    if(customInstructions){
+      systemPrompt = systemPrompt + "\n\n\n\n\n\nCustom Instructions:\n" + customInstructions;
+    }
+    // User message
+    let userMessage = "Documentation Content:\n"+docsContent;
+    // Optional context
+    let tokenEstimate = llm.estimateTokens(userMessage);
+    let consideredReferenceSourceCodeFiles = [];
+    if(referenceSourceCodeFiles?.length > 0){
+      userMessage += "\n\n\n\n\n\nReference Source Code:\n"+referenceFilesSummary+"\n\n";
+      for(const fileUrl of referenceSourceCodeFiles){
+        const fileContent = await FileUtility.readFile(fileUrl, this.accessMode);
+        tokenEstimate += llm.estimateTokens(fileContent);
+        if(tokenEstimate < llm.maxInputTokens){
+          consideredReferenceSourceCodeFiles.push(fileUrl);
+          userMessage += "// "+fileUrl+"\n\n"+fileContent+"\n\n\n\n\n\n";
+        }
+      }
+    }
+    if(context?.projectStructure){
+      tokenEstimate += llm.estimateTokens(context.projectStructure);
+      if(tokenEstimate < llm.maxInputTokens){
+        userMessage = userMessage + "\n\n\n\n\n\nProject Structure:\n"+context.projectStructure;
+      }
+    }
+    if(context?.relatedFileUrls){
+      // TODO: Remove duplicates
+      const relatedFilesContent = await Promise.all(context.relatedFileUrls.map(async (fileUrl) => {
+        return FileUtility.readFile(fileUrl, this.accessMode);
+      }));
+      tokenEstimate += llm.estimateTokens(relatedFilesContent.join("\n\n"));
+      if(tokenEstimate < llm.maxInputTokens){
+        userMessage = userMessage + "\n\n\n\n\n\nRelated Files Content:\n"+relatedFilesContent.join("\n\n");
+      }
+    }
+    if(context?.knowledgeBase){
+      tokenEstimate += llm.estimateTokens(context.knowledgeBase);
+      if(tokenEstimate < llm.maxInputTokens){
+        userMessage = userMessage + "\n\n\n\n\n\nKnowledge Base (Use this only for understanding. Do not include this in the reference docs, this is already published on the docs website.):\n"+context.knowledgeBase;
+      }
+    }
+    console.log("Auditing docs against code using", llm.model, "with temperature", llm.temperature, "and topP", llm.topP);
+    console.log("Total tokens estimated:", tokenEstimate);
+    console.log("Considered reference source code files:", consideredReferenceSourceCodeFiles?.length);
+    console.log("  ", consideredReferenceSourceCodeFiles?.map(fileUrl => fileUrl).join("\n   "));
+    if(!referenceSourceCodeFiles?.length){
+      return "Couldn't audit the docs against code because the reference source code files couldn't be found. Please make sure the repository URL is correct and the repository is accessible.";
+    } else if(!consideredReferenceSourceCodeFiles?.length){
+      return "Couldn't audit the docs against code because the reference source code files are too large to process. Please specify which specific code or class or configuration you want to audit.";
+    }
+    const auditReport = await llm.chat([
+      {
+        role: "system",
+        content: systemPrompt
+      },
+      {
+        role: "user",
+        content: [
+          {
+            type: "text",
+            text: userMessage
+          }
+        ]
+      }
+    ], null, {
+      isEnabled: true,
+      functionId: "auditDocsAgainstCode",
+      metadata: {
+        docsContentFilePath: docsContentFilePath,
+        referenceSourceCodeFiles: consideredReferenceSourceCodeFiles,
+        repoUrl: context?.repoUrl,
+        customInstructions: customInstructions,
+      }
+    });
+    console.log("Audit Report:", auditReport);
+    
+    // Add validation to ensure we always return a string
+    if (!auditReport || typeof auditReport !== 'string') {
+      console.error('LLM returned invalid response:', auditReport);
+      return "Error: Could not generate audit report. Please try again.";
+    }
+    
+    return auditReport + (auditReport?.length > 200 ? "\n\nNext Step: You may call editDocs tool to edit the docs to fix the issues found in the audit report" : "");
+  }
+
+  /**
+ * Generate reference docs for the given public interface and config files or GitHub repository
+ * @experimental This tool is experimental and might be dropped in the future.
+ * @param {string[]} referenceSourceCodeFiles - Remote URLs of the necessary source code files to generate reference docs e.g. [\"https://raw.githubusercontent.com/rudderlabs/analytics-go/refs/heads/master/analytics.go\", \"https://raw.githubusercontent.com/rudderlabs/analytics-go/refs/heads/master/config.go\"]
+ * @param {object} context - The context of the reference docs generation
+ * @param {string} context.interfaceType - Type of reference docs, how the code exposes its interfaces to the outside world
+ * @param {string} context.projectStructure - A nested markdown list of the project structure showing the folders/files hierarchy of the repository
+ * @param {string[]} context.relatedFileUrls - Remote URLs of the code files that are related to the given main public interface file
+ * @param {string} context.repoUrl - GitHub repository URL
+ * @param {string} context.knowledgeBase - Knowledge base such as higher level architecture of the parent project, this is only to develop understanding of the context of the current codebase. Do not include this in the reference docs, this is already published on docs website
+ * @param {string} customInstructions - Custom instructions for the reference docs generation process
+ * @returns {Promise<string>} - The generated reference docs
+ */
+  async generateReferenceDocs(referenceSourceCodeFiles, context, customInstructions){
+    console.log("Generating the docs", referenceSourceCodeFiles);
+    if(!referenceSourceCodeFiles?.length){
+      throw new Error("referenceSourceCodeFiles is required");
+    }
+    const llm = new LLM({
+      aiService: process.env.REVIEW_AI_SERVICE || "gemini",
+      model: process.env.REVIEW_AI_MODEL || "gemini-2.5-pro-preview-06-05"
+    });
+    // System prompt
+    let systemPrompt = GenerateReferenceDocsPrompt;
+    if(customInstructions){
+      systemPrompt = systemPrompt + "\n\n\n\n\n\nCustom Instructions:\n" + customInstructions;
+    }
+    let referenceSourceCodeFilesContent = "";
+    for(const fileUrl of referenceSourceCodeFiles){
+      const fileContent = await FileUtility.readFile(fileUrl, this.accessMode);
+      referenceSourceCodeFilesContent += "// "+fileUrl+"\n\n"+fileContent+"\n\n\n\n\n\n";
+    }
+    // User message
+    let userMessage = "Reference Source Code Files:\n"+referenceSourceCodeFilesContent;
+    if(context?.projectStructure){
+      userMessage = userMessage + "\n\n\n\n\n\nProject Structure:\n"+context.projectStructure;
+    }
+    // Optional context
+    let tokenEstimate = llm.estimateTokens(userMessage);
+    if(context?.relatedFileUrls){
+      // TODO: Remove duplicates
+      const relatedFilesContent = await Promise.all(context.relatedFileUrls.map(async (fileUrl) => {
+        return FileUtility.readFile(fileUrl, this.accessMode);
+      }));
+      tokenEstimate += llm.estimateTokens(relatedFilesContent.join("\n\n"));
+      if(tokenEstimate < llm.maxInputTokens){
+        userMessage = userMessage + "\n\n\n\n\n\nRelated Files Content:\n"+relatedFilesContent.join("\n\n");
+      }
+    }
+    if(context?.knowledgeBase){
+      tokenEstimate += llm.estimateTokens(context.knowledgeBase);
+      if(tokenEstimate < llm.maxInputTokens){
+        userMessage = userMessage + "\n\n\n\n\n\nKnowledge Base (Use this only for understanding. Do not include this in the reference docs, this is already published on the docs website.):\n"+context.knowledgeBase;
+      }
+    }
+    if(context?.repoUrl){
+      userMessage = userMessage + "\n\n\n\n\n\nSource Code Repository URL:\n"+context.repoUrl;
+    }
+    if(context?.interfaceType){
+      userMessage = userMessage + "\n\n\n\n\n\nInterface Type:\n"+context.interfaceType;
+    }
+    const generated = await llm.chat([
+      {
+        role: "system",
+        content: systemPrompt
+      },
+      {
+        role: "user",
+        content: [
+          {
+            type: "text",
+            text: userMessage
+          }
+        ]
+      }
+    ], null, {
+      isEnabled: true,
+      functionId: "generateReferenceDocs",
+      metadata: {
+        referenceSourceCodeFiles: referenceSourceCodeFiles,
+        repoUrl: context?.repoUrl,
+        interfaceType: context?.interfaceType,
+        customInstructions: customInstructions,
+      }
+    });
+    return generated;
+  }
+}
+
+export default DocsAgent;