Merge pull request #206 from olasunkanmi-SE/file-upload-feature

feat(file-upload): Enable local file uploads and processing

Author: olasunkanmi-SE

.gitignore

@@ -6,6 +6,7 @@ node_modules
 # Sentry Config File
 .sentryclirc
 config.ts
-patterns
+database
+files
 webviewUi/dist
 samples

src/agents/architecture.md

@@ -1,4 +1,4 @@
-```
+```mermaid
 graph TD
     A[GeminiLLM] -->|Initializes| B[GoogleGenerativeAI]
     A -->|Creates| C[Orchestrator]
@@ -36,4 +36,4 @@ graph TD
     Y[Function Call Retry] -->|Improves| L
 
     Z[Memory Management] -->|Optimizes| I
-```
+```

src/agents/file-upload.ts

@@ -6,19 +6,50 @@ import {
   GoogleGenAI,
 } from "@google/genai";
 import * as path from "path";
+import * as vscode from "vscode";
 import { BaseAiAgent } from "./base";
 import { Orchestrator } from "./orchestrator";
+import { IEventPayload } from "../emitter/interface";
 
-export class FileUploadAgent extends BaseAiAgent {
+export class FileUploadAgent extends BaseAiAgent implements vscode.Disposable {
   private readonly ai: GoogleGenAI;
   protected readonly orchestrator: Orchestrator;
   private static readonly PROCESSING_WAIT_TIME_MS = 6000;
   private static readonly MAX_CACHE_PAGE_SIZE = 10;
   private static readonly CACHE_MODEL = "gemini-1.5-flash-002";
+  private readonly disposables: vscode.Disposable[] = [];
+  private static instance: FileUploadAgent;
   constructor(private readonly apiKey: string) {
     super();
     this.ai = new GoogleGenAI({ apiKey: this.apiKey });
     this.orchestrator = Orchestrator.getInstance();
+    this.disposables.push(
+      this.orchestrator.onFileUpload(this.handleLocalFileUpload.bind(this)),
+    );
+  }
+
+  static initialize(apiKey: string) {
+    if (!FileUploadAgent.instance) {
+      FileUploadAgent.instance = new FileUploadAgent(apiKey);
+    }
+    return FileUploadAgent.instance;
+  }
+
+  private async handleLocalFileUpload(event: IEventPayload) {
+    try {
+      if (!event.message?.length) {
+        this.logger.info("Error while uploading file, try again");
+        return;
+      }
+      const data = JSON.parse(event.message);
+      const { fileName, filePath } = data;
+      if (fileName && filePath) {
+        await this.uploadAndProcessFile(filePath, fileName);
+      }
+    } catch (error: any) {
+      console.log(error);
+      throw new Error(error);
+    }
   }
 
   private async uploadFile(filePath: string, displayName: string) {
@@ -105,7 +136,7 @@ export class FileUploadAgent extends BaseAiAgent {
 
       return {
         response: response.text,
-        fileName,
+        fileName: fileName ?? file.displayName,
         cache: cached.name,
       };
     } catch (error) {
@@ -134,7 +165,8 @@ export class FileUploadAgent extends BaseAiAgent {
       model: FileUploadAgent.CACHE_MODEL,
       config: {
         contents: createUserContent(fileContent),
-        systemInstruction: "You are an expert analyzing documents",
+        systemInstruction:
+          "You are an expert analyzing documents. Give a detail analysis of this doc in markdown",
       },
     });
     this.logger.info("Cache created:", cached);
@@ -162,29 +194,7 @@ export class FileUploadAgent extends BaseAiAgent {
     }
   }
 
-  async getCaches(): Promise<CachedContent[]> {
-    this.logger.info("Retrieving caches");
-    const caches: CachedContent[] = [];
-
-    try {
-      const pager = await this.ai.caches.list({
-        config: { pageSize: FileUploadAgent.MAX_CACHE_PAGE_SIZE },
-      });
-
-      let page = pager.page;
-
-      do {
-        caches.push(...page);
-        if (!pager.hasNextPage()) break;
-        page = await pager.nextPage();
-      } while (true);
-
-      return caches;
-    } catch (error) {
-      console.log("Failed to retrieve caches", error);
-      throw new Error(
-        `Failed to retrieve caches: ${error instanceof Error ? error.message : String(error)}`,
-      );
-    }
+  public dispose(): void {
+    this.disposables.forEach((d) => d.dispose());
   }
 }

src/application/constant.ts

@@ -25,7 +25,7 @@ export enum COMMON {
 }
 export const GROQ_CONFIG = {
   temperature: 0.1,
-  max_tokens: 5024,
+  max_tokens: 500024,
   top_p: 1,
   stream: false,
   stop: null,

src/commands/knowledge-base.ts

@@ -13,12 +13,14 @@ export class EventEmitter extends BaseEmitter<Record<string, IEventPayload>> {
     this.createEvent("onSecretChange");
   onBootstrap: vscode.Event<IEventPayload> = this.createEvent("onBootstrap");
   onFileUpload: vscode.Event<IEventPayload> = this.createEvent("onFileUpload");
-  onFileProcessedSuccess: vscode.Event<IEventPayload> = this.createEvent(
+  onFileProcessSuccess: vscode.Event<IEventPayload> = this.createEvent(
     "onFileProcessSuccess",
   );
   onActiveworkspaceUpdate: vscode.Event<IEventPayload> = this.createEvent(
     "onActiveworkspaceUpdate",
   );
+  onFilesRetrieved: vscode.Event<IEventPayload> =
+    this.createEvent("onFilesRetrieved");
 
   /**
    * Emits a generic event with specified status, message, and optional data.

src/emitter/agent-emitter.ts

@@ -8,7 +8,9 @@ type AgentEventKeys =
   | "onSecretChange"
   | "onBootstrap"
   | "onActiveworkspaceUpdate"
-  | "onFileUpload";
+  | "onFileUpload"
+  | "onFileProcessSuccess"
+  | "onFilesRetrieved";
 
 export type IAgentEventMap = Record<AgentEventKeys, IEventPayload>;
 

src/emitter/interface.ts

@@ -14,7 +14,6 @@ import { GenerateCommitMessage } from "./commands/generate-commit-message";
 import { GenerateUnitTest } from "./commands/generate-unit-test";
 import { InLineChat } from "./commands/inline-chat";
 import { InterviewMe } from "./commands/interview-me";
-import { ReadFromKnowledgeBase } from "./commands/knowledge-base";
 import { OptimizeCode } from "./commands/optimize";
 import { RefactorCode } from "./commands/refactor";
 import { ReviewCode } from "./commands/review";
@@ -30,7 +29,9 @@ import { GroqWebViewProvider } from "./providers/groq";
 import { FileManager } from "./services/file-manager";
 import { initializeGenerativeAiEnvironment } from "./services/generative-ai-model-manager";
 import { Credentials } from "./services/github-authentication";
-import { getConfigValue } from "./utils/utils";
+import { getAPIKey, getConfigValue } from "./utils/utils";
+import { FileUploadAgent } from "./agents/file-upload";
+import * as fs from "fs";
 
 const {
   geminiKey,
@@ -50,9 +51,13 @@ const logger = new Logger("extension");
 let quickFixCodeAction: vscode.Disposable;
 let agentEventEmmitter: EventEmitter;
 
-async function connectToDatabase() {
+async function connectToDatabase(context: vscode.ExtensionContext) {
   try {
-    const dbPath = path.join(__dirname, "..", "patterns", "aii.db");
+    const dbDir = path.join(context.extensionPath, "database");
+    if (!fs.existsSync(dbDir)) {
+      fs.mkdirSync(dbDir);
+    }
+    const dbPath = path.join(__dirname, "..", "database", "aii.db");
     const urlPath = dbPath.replace(/\\/g, "/");
     const isWindows = dbPath.includes("\\");
     const filePath = isWindows ? `file:/${urlPath}` : `file:${urlPath}`;
@@ -65,7 +70,7 @@ async function connectToDatabase() {
 
 async function createFileDB(context: vscode.ExtensionContext) {
   try {
-    const fileUploader = new FileManager(context);
+    const fileUploader = new FileManager(context, "database");
     const files = await fileUploader.getFiles();
     if (!files?.find((file) => file.includes("dev.db"))) {
       await fileUploader.createFile("dev.db");
@@ -79,14 +84,16 @@ async function createFileDB(context: vscode.ExtensionContext) {
 export async function activate(context: vscode.ExtensionContext) {
   try {
     await createFileDB(context);
-    await connectToDatabase();
+    await connectToDatabase(context);
     const credentials = new Credentials();
+    const geminiApiKey = getAPIKey("gemini");
+    FileUploadAgent.initialize(geminiApiKey);
     await credentials.initialize(context);
     const session: vscode.AuthenticationSession | undefined =
       await credentials.getSession();
     logger.info(`Logged into GitHub as ${session?.account.label}`);
     Memory.getInstance();
-    const fileUpload = new FileManager(context);
+
     // TODO This is broken. Need to Fix
     // const index = CodeIndexingService.createInstance();
     // Get each of the folders and call the next line for each
@@ -136,11 +143,6 @@ export async function activate(context: vscode.ExtensionContext) {
       `${USER_MESSAGE} creates the code chart...`,
       context,
     );
-    const codePattern = fileUpload;
-    const knowledgeBase = new ReadFromKnowledgeBase(
-      `${USER_MESSAGE} generate your code pattern...`,
-      context,
-    );
     const generateCommitMessage = new GenerateCommitMessage(
       `${USER_MESSAGE} generates a commit message...`,
       context,
@@ -169,8 +171,6 @@ export async function activate(context: vscode.ExtensionContext) {
           errorMessage,
         ).execute(errorMessage),
       [explain]: async () => await explainCode.execute(),
-      [pattern]: async () => await codePattern.uploadFileHandler(),
-      [knowledge]: async () => await knowledgeBase.execute(),
       [commitMessage]: async () =>
         await generateCommitMessage.execute("commitMessage"),
       [generateCodeChart]: async () => await codeChartGenerator.execute(),

src/extension.ts

@@ -46,7 +46,7 @@ export class GeminiLLM
     this.logger = new Logger("GeminiLLM");
     this.groqLLM = GroqLLM.getInstance({
       apiKey: getAPIKey("groq"),
-      model: "deepseek-r1-distill-llama-70b",
+      model: "deepseek-r1-distill-qwen-32b",
     });
   }
 

src/llms/gemini/gemini.ts

@@ -31,7 +31,7 @@ export abstract class BaseWebViewProvider implements vscode.Disposable {
     protected readonly generativeAiModel: string,
     context: vscode.ExtensionContext,
   ) {
-    this.fileManager = FileManager.initialize(context);
+    this.fileManager = FileManager.initialize(context, "files");
     this.fileService = FileService.getInstance();
     this._context = context;
     this.orchestrator = Orchestrator.getInstance();
@@ -48,6 +48,7 @@ export abstract class BaseWebViewProvider implements vscode.Disposable {
       this.orchestrator.onActiveworkspaceUpdate(
         this.handleGenericEvents.bind(this),
       ),
+      this.orchestrator.onFileUpload(this.handleModelResponseEvent.bind(this)),
     );
   }
 
@@ -78,6 +79,7 @@ export abstract class BaseWebViewProvider implements vscode.Disposable {
     setTimeout(async () => {
       await this.publishWorkSpace();
     }, 2000);
+    await this.getFiles();
   }
 
   private async setWebviewHtml(view: vscode.WebviewView): Promise<void> {
@@ -87,6 +89,16 @@ export abstract class BaseWebViewProvider implements vscode.Disposable {
     );
   }
 
+  private async getFiles() {
+    const files: string[] = await this.fileManager.getFileNames();
+    if (files?.length) {
+      await this.currentWebView?.webview.postMessage({
+        type: "onFilesRetrieved",
+        message: JSON.stringify(files),
+      });
+    }
+  }
+
   private async publishWorkSpace(): Promise<void> {
     try {
       const filesAndDirs: IContextInfo =