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

File upload feature

Author: olasunkanmi-SE

package-lock.json

@@ -334,6 +334,7 @@
   },
   "dependencies": {
     "@anthropic-ai/sdk": "^0.20.9",
+    "@google/genai": "^0.7.0",
     "@google/generative-ai": "^0.21.0",
     "@googleapis/customsearch": "^3.2.0",
     "@libsql/client": "^0.14.0",
@@ -354,4 +355,4 @@
     "sinon": "^17.0.1"
   },
   "license": "MIT"
-}
+}

package.json

@@ -23,7 +23,7 @@ export abstract class BaseAiAgent
     };
   }
 
-  extractQueries(input: string): string[] | undefined {
+  protected extractQueries(input: string): string[] | undefined {
     const startIndex = input.indexOf("Queries: [");
     if (startIndex === -1) {
       return;
@@ -41,7 +41,7 @@ export abstract class BaseAiAgent
     return queriesArray;
   }
 
-  extractThought(input: string): string | undefined {
+  protected extractThought(input: string): string | undefined {
     const match = RegExp(/Thought:\s*(.*?)\n/).exec(input);
     if (match) {
       return match[1].trim();

src/agents/base.ts

@@ -0,0 +1,190 @@
+import {
+  CachedContent,
+  createPartFromUri,
+  createUserContent,
+  GenerateContentResponse,
+  GoogleGenAI,
+} from "@google/genai";
+import * as path from "path";
+import { BaseAiAgent } from "./base";
+import { Orchestrator } from "./orchestrator";
+
+export class FileUploadAgent extends BaseAiAgent {
+  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";
+  constructor(private readonly apiKey: string) {
+    super();
+    this.ai = new GoogleGenAI({ apiKey: this.apiKey });
+    this.orchestrator = Orchestrator.getInstance();
+  }
+
+  private async uploadFile(filePath: string, displayName: string) {
+    try {
+      return await this.ai.files.upload({
+        file: filePath,
+        config: { displayName },
+      });
+    } catch (error) {
+      console.error(`Failed to upload file: ${filePath}`, error);
+      throw new Error(
+        `File upload failed: ${error instanceof Error ? error.message : String(error)}`,
+      );
+    }
+  }
+
+  async uploadAndProcessFile(
+    filePath: string,
+    displayName: string,
+    prompt: string = "Summarize this document",
+  ): Promise<string | undefined> {
+    let file;
+    try {
+      file = (await this.uploadFile(filePath, displayName)) as any;
+      const processedFile = await this.waitForProcessing(file.name);
+
+      const result = await this.generateContentWithFile(processedFile, prompt);
+      if (result.response) {
+        this.orchestrator.publish("onResponse", JSON.stringify(result));
+      }
+      return result.response;
+    } catch (error) {
+      this.logger.info(`Failed to process file: ${file.name}`);
+      throw new Error(
+        `File processing pipeline failed: ${error instanceof Error ? error.message : String(error)}`,
+      );
+    }
+  }
+
+  private delay(ms: number): Promise<void> {
+    return new Promise((resolve) => setTimeout(resolve, ms));
+  }
+
+  private async waitForProcessing(fileName: string, maxRetries = 10) {
+    try {
+      let getFile = await this.ai.files.get({ name: fileName });
+      let retries = 0;
+      while (getFile.state === "PROCESSING" && retries < maxRetries) {
+        this.logger.info("☕ File upload in progress, grab a cup of coffee");
+        await this.delay(FileUploadAgent.PROCESSING_WAIT_TIME_MS);
+        getFile = await this.ai.files.get({ name: fileName });
+        retries++;
+      }
+      if (getFile.state === "FAILED") {
+        this.logger.info("File processing failed");
+      }
+      return getFile;
+    } catch (error: any) {
+      console.error("File processing failed.", error);
+      throw new Error(error.message);
+    }
+  }
+
+  private async generateContentWithFile(
+    file: any,
+    prompt: string,
+    cacheName?: string,
+  ): Promise<{
+    response: string | undefined;
+    fileName: string;
+    cache: string | undefined;
+  }> {
+    try {
+      const fileContent = createPartFromUri(file.uri, file.mimeType);
+      let cached = cacheName
+        ? await this.findOrCreateCache(cacheName, fileContent)
+        : await this.createNewCache(fileContent);
+
+      const response = await this.generateContentWithCache(
+        prompt,
+        cached.name ?? "",
+      );
+      const fileName = file.fsPath ? path.basename(file.fsPath) : "";
+
+      return {
+        response: response.text,
+        fileName,
+        cache: cached.name,
+      };
+    } catch (error) {
+      this.logger.error("Failed to generate content with file", error);
+      return {
+        response: undefined,
+        fileName: "",
+        cache: undefined,
+      };
+    }
+  }
+
+  private async findOrCreateCache(
+    cacheName: string,
+    fileContent: any,
+  ): Promise<CachedContent> {
+    try {
+      return await this.getDocCache(cacheName);
+    } catch (error) {
+      return await this.createNewCache(fileContent);
+    }
+  }
+
+  private async createNewCache(fileContent: any): Promise<CachedContent> {
+    const cached = await this.ai.caches.create({
+      model: FileUploadAgent.CACHE_MODEL,
+      config: {
+        contents: createUserContent(fileContent),
+        systemInstruction: "You are an expert analyzing documents",
+      },
+    });
+    this.logger.info("Cache created:", cached);
+    return cached;
+  }
+
+  private async generateContentWithCache(
+    prompt: string,
+    cacheName: string,
+  ): Promise<GenerateContentResponse> {
+    return await this.ai.models.generateContent({
+      model: FileUploadAgent.CACHE_MODEL,
+      contents: prompt,
+      config: { cachedContent: cacheName },
+    });
+  }
+
+  private async getDocCache(name: string): Promise<CachedContent> {
+    try {
+      const cache = await this.ai.caches.get({ name });
+      this.logger.info("Cache found:", cache);
+      return cache;
+    } catch (error: any) {
+      throw new Error("Cache not found or error occurred:", error);
+    }
+  }
+
+  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)}`,
+      );
+    }
+  }
+}

src/agents/file-upload.ts

@@ -349,3 +349,12 @@ export const PRIORITY_URLS = [
   "docs.oracle.com",
   "git-scm.com",
 ];
+
+const FILE_TYPE_PROMPTS: Record<string, string> = {
+  pdf: "Extract the main information and key points from this PDF document:",
+  docx: "Summarize the content of this Word document:",
+  csv: "Analyze this CSV data and provide key insights:",
+  txt: "Process this text file and extract the relevant information:",
+  json: "Parse this JSON file and provide a structured analysis:",
+  default: "Extract all relevant information from this file:",
+};

src/application/constant.ts

@@ -419,7 +419,7 @@ export abstract class EventGenerator implements IEventGenerator {
       }
     } catch (error) {
       this.logger.error(
-        "Error while passing bit response to the webview",
+        "Error while passing model response to the webview",
         error,
       );
     }

src/commands/event-generator.ts

@@ -12,6 +12,10 @@ export class EventEmitter extends BaseEmitter<Record<string, IEventPayload>> {
   onSecretChange: vscode.Event<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",
+  );
   onActiveworkspaceUpdate: vscode.Event<IEventPayload> = this.createEvent(
     "onActiveworkspaceUpdate",
   );

src/emitter/agent-emitter.ts

@@ -1,27 +1,16 @@
-interface IBaseEmitter {
-  timestamp: string;
-}
+type AgentEventKeys =
+  | "onStatus"
+  | "onError"
+  | "onUpdate"
+  | "onQuery"
+  | "onResponse"
+  | "onThinking"
+  | "onSecretChange"
+  | "onBootstrap"
+  | "onActiveworkspaceUpdate"
+  | "onFileUpload";
 
-export type Action = "query" | "update" | "status" | "error";
-export type EventState =
-  | "idle"
-  | "processing"
-  | "completed"
-  | "ui-update"
-  | "error"
-  | "query";
-
-export interface IAgentEventMap {
-  onStatus: IEventPayload;
-  onError: IEventPayload;
-  onUpdate: IEventPayload;
-  onQuery: IEventPayload;
-  onResponse: IEventPayload;
-  onThinking: IEventPayload;
-  onSecretChange: IEventPayload;
-  onBootstrap: IEventPayload;
-  onActiveworkspaceUpdate: IEventPayload;
-}
+export type IAgentEventMap = Record<AgentEventKeys, IEventPayload>;
 
 export interface IEventPayload {
   type: string;

src/emitter/interface.ts

@@ -87,6 +87,7 @@ export async function activate(context: vscode.ExtensionContext) {
     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
     // const result = await index.buildFunctionStructureMap();

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-qwen-32b",
+      model: "deepseek-r1-distill-llama-70b",
     });
   }