refactor: extract Map combination logic into fileMapUtils

src/tools/implementations/LintTool.ts

@@ -23,6 +23,7 @@ import { randomUUID } from 'crypto';
 import { tmpdir } from 'os';
 import { isLockFile } from '../../utils/diffGenerator.js';
 import { fileNameLookup } from '../../utils/fileNameLookup.js';
+import { getCombinedFileMap } from '../../utils/fileMapUtils.js';
 
 // --- Utility Functions ---
 
@@ -201,10 +202,7 @@ export const LintTool: MultiAgentTool = {
     }
 
     // Combine keys from both maps for a comprehensive lookup.
-    const allFilesMap = new Map<string, string>([
-      ...context.fileMap,
-      ...Array.from(context.binaryFileMap.keys()).map(key => [key, ''] as [string, string])
-    ]);
+    const allFilesMap = getCombinedFileMap(context.fileMap, context.binaryFileMap);
 
     let filename = providedFilename.trim();
 

src/tools/implementations/fileReaderTool.ts

@@ -16,6 +16,7 @@
 
 import { MultiAgentTool } from '../multiAgentTool.js';
 import { fileNameLookup } from '../../utils/fileNameLookup.js';
+import { getCombinedFileMap } from '../../utils/fileMapUtils.js';
 import { addDynamicallyRelevantFile, getFileAnalysis, updateFileEntry } from '../../utils/fileAnalysis.js';
 import { MultiAgentToolContext, MultiAgentToolResult, ToolParsingResult } from '../../momoa_core/types.js';
 import { getAssetString } from '../../services/promptManager.js';
@@ -46,10 +47,7 @@ export const fileReaderTool: MultiAgentTool = {
     }
 
     // Combine keys from both maps for a comprehensive lookup.
-    const allFilesMap = new Map<string, string>([
-      ...context.fileMap,
-      ...Array.from(context.binaryFileMap.keys()).map(key => [key, ''] as [string, string])
-    ]);
+    const allFilesMap = getCombinedFileMap(context.fileMap, context.binaryFileMap);
 
     // const filename = await fileNameLookup(providedFilename, allFilesMap, context.multiAgentGeminiClient);
     let filename = providedFilename.trim();

src/tools/implementations/fileSearchTool.ts

@@ -17,6 +17,7 @@
 import { MultiAgentTool } from '../multiAgentTool.js';
 import { findInFiles } from '../../utils/fileAnalysis.js';
 import { MultiAgentToolContext, MultiAgentToolResult, ToolParsingResult } from '../../momoa_core/types.js';
+import { getAllFileKeys } from '../../utils/fileMapUtils.js';
 
 /**
  * Implements the File Reader Tool, providing functionality to read file content
@@ -48,7 +49,7 @@ export const fileSearchTool: MultiAgentTool = {
     const searchResults = new Set<string>(contentMatches);
 
     // 2. Search all filenames (both text and binary).
-    const allFilenames = [...context.fileMap.keys(), ...context.binaryFileMap.keys()];
+    const allFilenames = getAllFileKeys(context.fileMap, context.binaryFileMap);
     for (const filename of allFilenames) {
         if (filename.includes(query)) {
             if (context.binaryFileMap.has(filename)) {

src/tools/implementations/renameFolderTool.ts

@@ -14,8 +14,9 @@
  * limitations under the License.
  */
 
-import { MultiAgentToolContext, MultiAgentToolResult, ToolParsingResult } from "../../momoa_core/types.js";
-import { addDynamicallyRelevantFile, getFileAnalysis, removeFileEntry, updateFileEntry } from "../../utils/fileAnalysis.js";
+import { MultiAgentToolContext, MultiAgentToolResult, ToolParsingResult } from '../../momoa_core/types.js';
+import { addDynamicallyRelevantFile, getFileAnalysis, removeFileEntry, updateFileEntry } from '../../utils/fileAnalysis.js';
+import { getAllFileKeys } from '../../utils/fileMapUtils.js';
 import { MultiAgentTool } from "../multiAgentTool.js";
 
 /**
@@ -42,7 +43,7 @@ export const moveFolderTool: MultiAgentTool = {
     addDynamicallyRelevantFile(destination);
 
     let sourceExists = false;
-    const allFileKeys = [...context.fileMap.keys(), ...context.binaryFileMap.keys()];
+    const allFileKeys = getAllFileKeys(context.fileMap, context.binaryFileMap);
 
     // Check if the source path exists as a direct file entry
     if (context.fileMap.has(source) || context.binaryFileMap.has(source)) {

src/tools/implementations/revertFileTool.ts

@@ -17,6 +17,7 @@
 import { MultiAgentTool } from '../multiAgentTool.js';
 import { MultiAgentToolContext, MultiAgentToolResult, ToolParsingResult } from '../../momoa_core/types.js';
 import { fileNameLookup } from '../../utils/fileNameLookup.js';
+import { getCombinedFileMap } from '../../utils/fileMapUtils.js';
 import { generateDiffString, isLockFile, getLockFileHiddenPlaceholder } from '../../utils/diffGenerator.js';
 import { getAssetString } from '../../services/promptManager.js';
 import { updateFileEntry } from '../../utils/fileAnalysis.js';
@@ -59,10 +60,7 @@ export const revertFileTool: MultiAgentTool = {
     }
 
     // Use fileNameLookup to find the precise file name
-    const allFilesMap = new Map<string, string>([
-      ...context.fileMap,
-      ...Array.from(context.binaryFileMap.keys()).map(key => [key, ''] as [string, string])
-    ]);
+    const allFilesMap = getCombinedFileMap(context.fileMap, context.binaryFileMap);
     const filename = await fileNameLookup(providedFilename, allFilesMap, context.multiAgentGeminiClient);
 
     const originalExisted = context.originalFileMap.has(filename) || context.originalBinaryFileMap.has(filename);

src/tools/implementations/smartFileEditorTool.ts

@@ -16,6 +16,7 @@
 
 import { MultiAgentTool } from '../multiAgentTool.js';
 import { fileNameLookup } from '../../utils/fileNameLookup.js';
+import { getCombinedFileMap } from '../../utils/fileMapUtils.js';
 import { fileReaderTool } from './fileReaderTool.js';
 import { addDynamicallyRelevantFile, removeFileEntry, updateFileEntry } from '../../utils/fileAnalysis.js';
 import { FileOperation, MultiAgentToolContext, MultiAgentToolResult, ToolParsingResult } from '../../momoa_core/types.js';
@@ -442,10 +443,7 @@ export const smartFileEditorTool: MultiAgentTool = {
     }
 
     // Combine all file keys for a comprehensive lookup.
-    const allFilesMap = new Map<string, string>([
-        ...context.fileMap,
-        ...Array.from(context.binaryFileMap.keys()).map(key => [key, ''] as [string, string])
-    ]);
+    const allFilesMap = getCombinedFileMap(context.fileMap, context.binaryFileMap);
 
     if (isUnambiguousCreate) {
         // This is a CREATE or OVERWRITE. The file does not need to exist.

src/utils/diffGenerator.ts

@@ -17,6 +17,7 @@
 import { applyPatch, createTwoFilesPatch, parsePatch } from 'diff';
 import { MultiAgentToolContext } from '../momoa_core/types.js';
 import { Buffer } from 'node:buffer';
+import { getCombinedFileMap } from './fileMapUtils.js';
 import * as zlib from 'node:zlib';
 
 const LOCK_FILE_NAMES = new Set([
@@ -268,14 +269,12 @@ export function applyDiff(
 
 export function generateDiffString(context: MultiAgentToolContext, redactFiles: boolean = false): string {
   const binaryPlaceholder = "[binary file content]";
-  const combinedOriginalFileMap = new Map<string, string>([
-      ...(context.originalFileMap || new Map()),
-      ...Array.from(context.originalBinaryFileMap?.keys() || []).map(key => [key, binaryPlaceholder] as [string, string])
-  ]);
-  const combinedFinalFileMap = new Map<string, string>([
-      ...context.fileMap,
-      ...Array.from(context.binaryFileMap.keys()).map(key => [key, binaryPlaceholder] as [string, string])
-  ]);
+  const combinedOriginalFileMap = getCombinedFileMap(
+      context.originalFileMap || new Map(),
+      context.originalBinaryFileMap || new Map(),
+      binaryPlaceholder
+  );
+  const combinedFinalFileMap = getCombinedFileMap(context.fileMap, context.binaryFileMap, binaryPlaceholder);
 
   const allBinaryFiles = new Set([
       ...(context.originalBinaryFileMap?.keys() || []),

src/utils/fileMapUtils.ts

@@ -0,0 +1,49 @@
+/**
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/**
+ * Combines a file map and a binary file map into a single map.
+ * Binary files are assigned a placeholder value.
+ *
+ * @param fileMap The map of text files.
+ * @param binaryFileMap The map of binary files.
+ * @param binaryPlaceholder The placeholder value to use for binary files.
+ * @returns A new map containing all files.
+ */
+export function getCombinedFileMap(
+  fileMap: Map<string, string>,
+  binaryFileMap: Map<string, string>,
+  binaryPlaceholder: string = ''
+): Map<string, string> {
+  return new Map<string, string>([
+    ...fileMap,
+    ...Array.from(binaryFileMap.keys()).map(key => [key, binaryPlaceholder] as [string, string])
+  ]);
+}
+
+/**
+ * Returns an array of all file keys from both the file map and the binary file map.
+ *
+ * @param fileMap The map of text files.
+ * @param binaryFileMap The map of binary files.
+ * @returns An array of all filenames.
+ */
+export function getAllFileKeys(
+  fileMap: Map<string, string>,
+  binaryFileMap: Map<string, string>
+): string[] {
+  return [...fileMap.keys(), ...binaryFileMap.keys()];
+}

src/utils/markdownUtils.ts

@@ -18,6 +18,7 @@ import { MultiAgentToolContext } from "../momoa_core/types.js";
 import { getAssetString } from "../services/promptManager.js";
 import { getFileAnalysis } from "./fileAnalysis.js";
 import { fileNameLookup } from "./fileNameLookup.js";
+import { getCombinedFileMap } from "./fileMapUtils.js";
 
 /**
  * Removes triple backtick fences from a string if they exist at the start and end.
@@ -211,10 +212,7 @@ export async function getFilesAndContent(requestedFiles: {FILENAME: string, DESC
   let result = '--No Files--';
   if (requestedFiles && requestedFiles.length > 0 && (context.fileMap || context.binaryFileMap)) {
     // Combine text and binary file maps for a comprehensive file lookup.
-    const allFilesMap = new Map<string, string>([
-      ...context.fileMap,
-      ...Array.from(context.binaryFileMap.keys()).map(key => [key, ''] as [string, string])
-    ]);
+    const allFilesMap = getCombinedFileMap(context.fileMap, context.binaryFileMap);
 
     const fileResults: string[] = [];