Merge pull request #161 from olasunkanmi-SE/refactor-code-indexing

olasunkanmi-SE · web-flow · commit 593cdb1a4960 · 2025-02-28T22:13:34.000+08:00
Refactor(code-indexing): Improve code indexing and context retrieval
diff --git a/.github/workflows/workflow.yml b/.github/workflows/workflow.yml
@@ -1,4 +1,4 @@
-name: CI for restaurant
+name: CI Pipeline
 
 on:
   pull_request:
diff --git a/src/agents/orchestrator.ts b/src/agents/orchestrator.ts
@@ -11,7 +11,7 @@ export class Orchestrator extends BaseAiAgent implements vscode.Disposable {
     this.disposables.push(
       this.onStatusChange(this.handleStatus.bind(this)),
       this.onPromptGenerated(this.handlePromptGeneratedEvent.bind(this)),
-      this.onError(this.handleError.bind(this)),
+      // this.onError(this.handleError.bind(this)),
     );
   }
 
@@ -31,9 +31,9 @@ export class Orchestrator extends BaseAiAgent implements vscode.Disposable {
     this.publish("onResponse", event.message);
   }
 
-  public handleError(event: IEventPayload) {
-    console.log(`Error: ${event.message}`);
-  }
+  // public handleError(event: IEventPayload) {
+  //   console.log(`Error: ${event.message}`);
+  // }
 
   public dispose(): void {
     this.disposables.forEach((d) => d.dispose());
diff --git a/src/application/constant.ts b/src/application/constant.ts
@@ -65,6 +65,7 @@ export enum FSPROPS {
   SRC_DIRECTORY = "src",
   TS_FILE_PATTERN = "**/*.ts",
   TSCONFIG_FILE = "tsconfig.json",
+  NODE_MODULES_PATTERN = "**/node_modules/**",
 }
 
 export const EmbeddingsConfig = {
diff --git a/src/application/interfaces/ts.code.mapper.interface.ts b/src/application/interfaces/ts.code.mapper.interface.ts
@@ -105,7 +105,7 @@ export interface ITypeScriptCodeMapper {
    * Builds a hierarchical map of the codebase by traversing TypeScript files
    * and extracting module and class information.
    */
-  buildCodebaseMap(): Promise<ICodebaseMap>;
+  buildCodebaseMap(): Promise<Map<string, string>>;
 
   /**
    * Extracts interface information from a TypeScript interface declaration.
diff --git a/src/extension.ts b/src/extension.ts
@@ -44,7 +44,7 @@ const {
 
 const connectDB = async () => {
   await dbManager.connect(
-    "file:/Users/olasunkanmioyinlola/Documents/Apps/codebuddy/patterns/aii.db",
+    "file:/Users/olasunkanmi/Documents/Github/codebuddy/patterns/dev.db",
   );
 };
 
@@ -76,10 +76,11 @@ export async function activate(context: vscode.ExtensionContext) {
     // const names = await fileUpload.getFileNames();
     // console.log(files, names);
 
-    // const index = CodeIndexingService.createInstance();
-    // const result = await index.buildFunctionStructureMap();
+    const index = CodeIndexingService.createInstance();
+    // Get each of the folders and call the next line for each
+    const result = await index.buildFunctionStructureMap();
     // await index.insertFunctionsinDB();
-    // console.log(result);
+    console.log(result);
     const {
       comment,
       review,
diff --git a/src/llms/gemini/gemini.ts b/src/llms/gemini/gemini.ts
@@ -280,6 +280,7 @@ export class GeminiLLM
 
       return finalResult;
     } catch (error: any) {
+      // Note. Use this approach to return error messages back to the FE
       this.orchestrator.publish(
         "onError",
         "Model not responding at this time, please try again",
diff --git a/src/providers/base.ts b/src/providers/base.ts
@@ -103,11 +103,6 @@ export abstract class BaseWebViewProvider implements vscode.Disposable {
   public handleModelResponseEvent(event: IEventPayload) {
     this.sendResponse(formatText(event.message), "bot");
   }
-
-  public handleThinkingEvent(event: IEventPayload) {
-    this.sendResponse(formatText(event.message), "bot");
-  }
-
   abstract generateResponse(
     message?: string,
     metaData?: Record<string, any>,
diff --git a/src/services/code-indexing.ts b/src/services/code-indexing.ts
@@ -47,6 +47,7 @@ export class CodeIndexingService {
    */
   async buildFunctionStructureMap(): Promise<Partial<IFunctionData>[]> {
     try {
+      // TODO Get all the typescript project compilers information
       const codeATS = await TypeScriptAtsMapper.getInstance();
       if (!codeATS) {
         throw new Error("Failed to get TypeScriptAtsMapper instance");
diff --git a/src/services/context-retriever.ts b/src/services/context-retriever.ts
@@ -9,6 +9,7 @@ import { CodeRepository } from "../infrastructure/repository/code";
 import { getGeminiAPIKey } from "../utils/utils";
 import { EmbeddingService } from "./embedding";
 import { WebSearchService } from "./web-search-service";
+import { Orchestrator } from "../agents/orchestrator";
 
 export class ContextRetriever {
   private readonly codeRepository: CodeRepository;
@@ -17,12 +18,14 @@ export class ContextRetriever {
   private readonly logger: Logger;
   private static instance: ContextRetriever;
   private readonly webSearchService: WebSearchService;
+  protected readonly orchestrator: Orchestrator;
   constructor() {
     this.codeRepository = CodeRepository.getInstance();
     const geminiApiKey = getGeminiAPIKey();
     this.embeddingService = new EmbeddingService(geminiApiKey);
     this.logger = new Logger("ContextRetriever");
     this.webSearchService = WebSearchService.getInstance();
+    this.orchestrator = Orchestrator.getInstance();
   }
 
   static initialize() {
@@ -36,6 +39,7 @@ export class ContextRetriever {
     try {
       const embedding = await this.embeddingService.generateEmbedding(input);
       this.logger.info("Retrieving context from DB");
+      // this.orchestrator.publish("onUpdate", "Retrieving context from DB");
       return this.codeRepository.searchSimilarFunctions(
         embedding,
         ContextRetriever.SEARCH_RESULT_COUNT,
@@ -80,8 +84,8 @@ export class ContextRetriever {
 
   async webSearch(query: string) {
     try {
-      const x = Array.isArray(query) ? query.join("") : query;
-      return await this.webSearchService.run(x);
+      const text = Array.isArray(query) ? query.join("") : query;
+      return await this.webSearchService.run(text);
     } catch (error) {
       this.logger.error("Error reading file:", error);
       throw error;
diff --git a/src/services/file-system.ts b/src/services/file-system.ts
@@ -6,9 +6,18 @@ import { Logger } from "../infrastructure/logger/logger";
 
 export class FileService {
   readonly logger: Logger;
+  private static instance: FileService;
   constructor() {
     this.logger = new Logger("FileSystemService");
   }
+
+  static getInstance() {
+    if (!FileService.instance) {
+      FileService.instance = new FileService();
+    }
+    return FileService.instance;
+  }
+
   getWorkspaceInfo(dir: string): IWorkspaceInfo | undefined {
     try {
       const workspaceFolder = vscode.workspace.workspaceFolders?.[0];
@@ -35,25 +44,29 @@ export class FileService {
       const directories = await vscode.workspace.fs.readDirectory(
         workSpaceInfo.root,
       );
+      const tsFilePaths: string[] = [];
+      for (const [name, type] of directories) {
+        if (type === vscode.FileType.Directory) {
+          const dirUri = vscode.Uri.joinPath(workSpaceInfo.root, name);
+          const tsConfigUrl = vscode.Uri.joinPath(
+            dirUri,
+            FSPROPS.TSCONFIG_FILE,
+          );
 
-      const directory = directories.filter(
-        ([name, type]) => type === vscode.FileType.Directory && name === dir,
-      );
-
-      if (!directory) {
-        throw Error(`${dir} does not exist within this workspace`);
+          const tsConfigExists =
+            (await Promise.resolve(vscode.workspace.fs.stat(tsConfigUrl)).catch(
+              () => null,
+            )) !== null;
+          if (tsConfigExists) {
+            const tsFiles = await vscode.workspace.findFiles(
+              new vscode.RelativePattern(dirUri, pattern),
+              FSPROPS.NODE_MODULES_PATTERN,
+            );
+            tsFilePaths.push(...tsFiles.map((file) => file.fsPath));
+          }
+        }
       }
-
-      const directoryFiles = directory.map(async ([file]) => {
-        const srcUri = vscode.Uri.joinPath(workSpaceInfo.root, file);
-        const srcFiles = await vscode.workspace.findFiles(
-          new vscode.RelativePattern(srcUri, pattern),
-        );
-        return srcFiles.map((file) => file.fsPath);
-      });
-
-      const srcFilePaths = await Promise.all(directoryFiles);
-      return srcFilePaths.flat();
+      return tsFilePaths;
     } catch (error) {
       handleError(
         error,
diff --git a/src/services/typescript-ats.service.ts b/src/services/typescript-ats.service.ts
@@ -21,12 +21,10 @@ import { FileService } from "./file-system";
 export class TypeScriptAtsMapper implements ITypeScriptCodeMapper {
   private program: ts.Program | undefined;
   private typeChecker: ts.TypeChecker | undefined;
-  private fsService: FileService | undefined;
+  private readonly fsService: FileService;
   private static instance: TypeScriptAtsMapper;
   constructor() {
-    if (!this.fsService) {
-      this.fileSysService();
-    }
+    this.fsService = FileService.getInstance();
   }
 
   public async init() {
@@ -36,15 +34,12 @@ export class TypeScriptAtsMapper implements ITypeScriptCodeMapper {
   public static async getInstance(): Promise<TypeScriptAtsMapper> {
     if (!TypeScriptAtsMapper.instance) {
       TypeScriptAtsMapper.instance = new TypeScriptAtsMapper();
+      // TODO remove this init flow at this point
       await TypeScriptAtsMapper.instance.init();
     }
     return TypeScriptAtsMapper.instance;
   }
 
-  fileSysService() {
-    return (this.fsService = new FileService());
-  }
-
   /**
    * Initializes a TypeScript program by reading the TS configuration file and creating a new program instance.
    * This method sets up the program and type checker for further compilation and analysis.
@@ -53,10 +48,7 @@ export class TypeScriptAtsMapper implements ITypeScriptCodeMapper {
   // Users may have mono repos or may open the folder at an higher level folder
   private async initializeTypescriptProgram(): Promise<void> {
     try {
-      if (!this.fsService) {
-        this.fsService = new FileService();
-      }
-      const fileContent = await this.fsService?.readFile(FSPROPS.TSCONFIG_FILE);
+      const fileContent = await this.fsService.readFile(FSPROPS.TSCONFIG_FILE);
       if (!fileContent) {
         throw Error;
       }
@@ -124,12 +116,13 @@ export class TypeScriptAtsMapper implements ITypeScriptCodeMapper {
     node: DeclarationFunctionNode,
     sourceFile: ts.SourceFile,
     info: IClassInfo | IModuleInfo,
-  ): void {
+  ): IFunctionInfo | null {
     const functionInfo: IFunctionInfo | null =
       this.getFunctionDetails(node, sourceFile) ?? null;
     if (functionInfo) {
       info?.functions?.push(functionInfo);
     }
+    return functionInfo;
   }
 
   /**
@@ -162,6 +155,7 @@ export class TypeScriptAtsMapper implements ITypeScriptCodeMapper {
     if (interfaceInfo) {
       info?.interfaces?.push(interfaceInfo);
     }
+    return interfaceInfo;
   }
 
   /**
@@ -195,15 +189,20 @@ export class TypeScriptAtsMapper implements ITypeScriptCodeMapper {
     sourceFile: ts.SourceFile,
     info: IClassInfo | IModuleInfo,
     member?: ts.ClassElement,
-  ): void {
+  ): IFunctionInfo | null {
     try {
-      const currentElement = member ? member : node;
+      let classFunctions = null;
+      const currentElement = member ?? node;
       if (
         ts.isMethodDeclaration(currentElement) ||
         ts.isFunctionDeclaration(currentElement) ||
         ts.isArrowFunction(currentElement)
       ) {
-        this.aggregateFunctions(currentElement, sourceFile, info);
+        classFunctions = this.aggregateFunctions(
+          currentElement,
+          sourceFile,
+          info,
+        );
       }
 
       if (ts.isPropertyDeclaration(currentElement)) {
@@ -217,6 +216,7 @@ export class TypeScriptAtsMapper implements ITypeScriptCodeMapper {
       if (ts.isEnumDeclaration(node)) {
         this.aggregateEnums(node, sourceFile, info);
       }
+      return classFunctions;
     } catch (error) {
       handleError(error, `Unable to process class members`);
       throw error;
@@ -486,19 +486,19 @@ export class TypeScriptAtsMapper implements ITypeScriptCodeMapper {
    * Builds a hierarchical map of the codebase by traversing TypeScript files
    * and extracting module and class information.
    */
-  async buildCodebaseMap(): Promise<ICodebaseMap> {
+  async buildCodebaseMap(): Promise<Map<string, string>> {
     try {
-      const rootDir: string = process.cwd();
+      const filesMap = new Map<string, string>();
+      const rootDir: string = this.fsService.getRootFilePath();
       const normalizedRootDir = path.normalize(rootDir);
       const codebaseMap: ICodebaseMap = {};
       const repoNames: string = path.basename(normalizedRootDir);
       codebaseMap[repoNames] = { modules: {} };
 
-      const tsFiles: string[] | undefined =
-        await this.fsService?.getFilesFromDirectory(
-          FSPROPS.SRC_DIRECTORY,
-          FSPROPS.TS_FILE_PATTERN,
-        );
+      const tsFiles = await this.fsService.getFilesFromDirectory(
+        rootDir,
+        FSPROPS.TS_FILE_PATTERN,
+      );
       if (!tsFiles?.length) {
         throw Error(`No Typescript files found`);
       }
@@ -516,27 +516,53 @@ export class TypeScriptAtsMapper implements ITypeScriptCodeMapper {
         );
         ts.forEachChild(sourceFile, (node) => {
           if (ts.isClassDeclaration(node)) {
-            const classInfo = this.extractClassMetaData(node, sourceFile);
+            const classInfo: IClassInfo = this.extractClassMetaData(
+              node,
+              sourceFile,
+            );
+
             if (classInfo) {
+              if (classInfo.name) {
+                filesMap.set(classInfo.name, moduleRalativePath);
+              }
               moduleInfo?.classes?.push(classInfo);
             }
-            this.processClassMembers(node, sourceFile, moduleInfo);
+            const functionInfo: IFunctionInfo | null = this.processClassMembers(
+              node,
+              sourceFile,
+              moduleInfo,
+            );
+            if (functionInfo?.name) {
+              filesMap.set(functionInfo.name, moduleRalativePath);
+            }
           }
 
           if (
             ts.isMethodDeclaration(node) ||
             ts.isFunctionDeclaration(node) ||
             (ts.isVariableDeclaration(node) && ts.isArrowFunction(node))
           ) {
-            this.aggregateFunctions(node, sourceFile, moduleInfo);
+            const functionInfo = this.aggregateFunctions(
+              node,
+              sourceFile,
+              moduleInfo,
+            );
+            if (functionInfo?.name) {
+              filesMap.set(functionInfo.name, moduleRalativePath);
+            }
           }
 
           if (ts.isPropertyDeclaration(node)) {
             this.aggergateProperties(node, sourceFile, moduleInfo);
           }
 
           if (ts.isInterfaceDeclaration(node)) {
-            this.aggregateInterfaces(node, sourceFile, moduleInfo);
+            const interfaceInfo = this.aggregateInterfaces(
+              node,
+              sourceFile,
+              moduleInfo,
+            );
+            filesMap.set(interfaceInfo.name, moduleRalativePath);
           }
 
           if (ts.isEnumDeclaration(node)) {
@@ -545,7 +571,8 @@ export class TypeScriptAtsMapper implements ITypeScriptCodeMapper {
           codebaseMap[repoNames].modules[moduleRalativePath] = moduleInfo;
         });
       });
-      return codebaseMap;
+      //For Code Indexing return codebaseMap
+      return filesMap;
     } catch (error) {
       handleError(error, "Error fetching the files");
       throw Error;
diff --git a/webviewUi/src/components/botMessage.tsx b/webviewUi/src/components/botMessage.tsx

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-name: CI for restaurant`
	`1`	`+name: CI Pipeline`
`2`	`2`
`3`	`3`	`on:`
`4`	`4`	`pull_request:`
Original file line number	Diff line number	Diff line change
`@@ -65,6 +65,7 @@ export enum FSPROPS {`
`65`	`65`	`SRC_DIRECTORY = "src",`
`66`	`66`	`TS_FILE_PATTERN = "*/.ts",`
`67`	`67`	`TSCONFIG_FILE = "tsconfig.json",`
	`68`	`+ NODE_MODULES_PATTERN = "/node_modules/",`
`68`	`69`	`}`
`69`	`70`
`70`	`71`	`export const EmbeddingsConfig = {`