Skip to content

Commit

Permalink
fix: Migrates gemini cache from os temp dir to workspace cache (#219)
Browse files Browse the repository at this point in the history 
* fix: Migrates gemini cache from os temp dir to workspace cache

* Minor changes

* New cache structure for gemini cache

* fix: chache stored seperatly for each project

---------

Co-authored-by: Samyak Jain <[email protected]>
  • Loading branch information
@KevalPrajapati @samyakkkk
KevalPrajapati and samyakkkk authored Mar 21, 2024
1 parent e8c0509 commit 6b01c47
Show file tree
Hide file tree
Showing 6 changed files with 126 additions and 36 deletions.
1 change: 0 additions & 1 deletion vscode/src/extension.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
// Import the module and reference it with the alias vscode in your code below
import * as vscode from 'vscode';

import { createInlineCodeCompletion } from './tools/create/inline_code_completion';
import { makeHttpRequest } from './repository/http-utils';
import { activateTelemetry, logEvent } from './utilities/telemetry-reporter';
import * as dotenv from 'dotenv';
Expand Down
69 changes: 38 additions & 31 deletions vscode/src/repository/gemini-repository.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ import { getReferenceEditor } from "../utilities/state-objects";
import { GenerationRepository } from "./generation-repository";
import { extractReferenceTextFromEditor } from "../utilities/code-processing";
import { logError } from "../utilities/telemetry-reporter";
import { CacheManager } from "../utilities/cache-manager";
import { computeCodehash } from "../shared/utils";

function handleError(error: Error, userFriendlyMessage: string): never {
console.error(error);
Expand All @@ -22,12 +24,37 @@ export class GeminiRepository extends GenerationRepository {
constructor(apiKey: string) {
super(apiKey);
this.genAI = new GoogleGenerativeAI(apiKey);
this.ensureCacheDirExists().catch(error => {
handleError(error, 'Failed to initialize the cache directory.');
});
this.migrateCache();
GeminiRepository._instance = this;
}

/**
* Migrates the cache from the operating system's temporary directory to the extension's cache directory.
*
* The cache was previously stored in the operating system's temporary directory, but it is now being moved to the extension's cache directory.
* This migration ensures that the cache is stored in a secure location and avoids potential issues with file permissions.
*
* This is a one-time operation and will be removed in future releases.
*/
private async migrateCache() {
try {
// TODO: Need to update the migration logic for the new cache structure
// Check if the cache file exists in the old location
if (fs.existsSync(this.cacheFilePath)) {
// Move the cache file to the CacheManager
const cacheData = await fs.promises.readFile(this.cacheFilePath, 'utf8');
console.log(cacheData);
if (cacheData.length > 0) {
await CacheManager.getInstance().setGeminiCache(this.codehashCache);
}
// Delete the old cache file
await fs.promises.unlink(this.cacheFilePath);
}
} catch (error) {
console.error('Failed to migrate the cache:', error);
}
}

public static getInstance(): GeminiRepository {
return GeminiRepository._instance;
}
Expand Down Expand Up @@ -124,7 +151,7 @@ export class GeminiRepository extends GenerationRepository {
throw new Error('No workspace folders found.');
}
const projectFolder = workspaceFolders[0].uri.fsPath; // Assuming single root workspace
const hash = this.computeCodehash(projectFolder); // Hash the path for uniqueness
const hash = computeCodehash(projectFolder); // Hash the path for uniqueness
// Use os.tmpdir() to get the system's temporary directory
const tempDir = require('os').tmpdir();
return require('path').join(tempDir, 'flutterGPT', `${hash}.codehashCache.json`);
Expand All @@ -135,8 +162,7 @@ export class GeminiRepository extends GenerationRepository {
private async saveCache() {
try {
const cacheData = JSON.stringify(this.codehashCache);
const cachePath = this.cacheFilePath;
await fs.promises.writeFile(cachePath, cacheData, { encoding: 'utf8', mode: 0o600 }); // Sets the file mode to read/write for the owner only
await CacheManager.getInstance().setGeminiCache(this.codehashCache);
} catch (error) {
if (error instanceof Error) {
handleError(error, 'Failed to save the cache data.');
Expand All @@ -148,37 +174,18 @@ export class GeminiRepository extends GenerationRepository {

private async loadCache() {
try {
if (fs.existsSync(this.cacheFilePath)) {
const cacheData = await fs.promises.readFile(this.cacheFilePath, 'utf8');
const cacheData = await CacheManager.getInstance().getGeminiCache();
if (cacheData) {
this.codehashCache = JSON.parse(cacheData);
}
else {
this.codehashCache = {};
}
} catch (error) {
console.error("Error loading cache: ", error);
}
}

// Ensure the directory exists and has the correct permissions
private async ensureCacheDirExists() {
const cacheDir = path.dirname(this.cacheFilePath);
try {
// Create the directory if it doesn't exist
if (!fs.existsSync(cacheDir)) {
await fs.promises.mkdir(cacheDir, { recursive: true, mode: 0o700 }); // Sets the directory mode to read/write/execute for the owner only
}
} catch (error: any) {
if (error.code !== 'EEXIST') {
handleError(error as Error, 'Failed to create a secure cache directory.');
}
}
}

// Compute a codehash for file contents
private computeCodehash(fileContents: string): string {
// Normalize the file content by removing whitespace and newlines
const normalizedContent = fileContents.replace(/\s+/g, '');
return crypto.createHash('sha256').update(normalizedContent).digest('hex');
}

// Find 5 closest dart files for query
public async findClosestDartFiles(query: string, view?: vscode.WebviewView, shortcut: boolean = false, filepath: string = ''): Promise<string> {
//start timer
Expand Down Expand Up @@ -219,7 +226,7 @@ export class GeminiRepository extends GenerationRepository {
const document = await vscode.workspace.openTextDocument(file);
const relativePath = vscode.workspace.asRelativePath(file, false);
const text = `File name: ${file.path.split('/').pop()}\nFile path: ${relativePath}\nFile code:\n\n\`\`\`dart\n${document.getText()}\`\`\`\n\n------\n\n`;
const codehash = this.computeCodehash(text);
const codehash = computeCodehash(text);
return {
text,
path: file.path,
Expand Down
9 changes: 8 additions & 1 deletion vscode/src/shared/utils.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import { GeminiRepository } from '../repository/gemini-repository';
import { ILspAnalyzer } from './types/LspAnalyzer';
import { Outline } from './types/custom_protocols';
import { SemanticTokensRegistrationType, SemanticTokensProviderShape, SymbolKind } from 'vscode-languageclient';
import * as crypto from 'crypto';

export async function getCodeForElementAtRange(analyzer: ILspAnalyzer, document: vscode.TextDocument, range: vscode.Range): Promise<string | undefined> {
const outline = (await analyzer.fileTracker.waitForOutline(document));
Expand Down Expand Up @@ -46,7 +47,7 @@ export async function getCodeForElementAtRange(analyzer: ILspAnalyzer, document:
}

export async function cursorIsAt(type: String, analyzer: ILspAnalyzer, document: vscode.TextDocument, activeTextEditor: vscode.TextEditor | undefined, range: vscode.Range, strict: boolean = true): Promise<{ symbolRange: vscode.Range, symbol: Outline } | undefined> {
if (!activeTextEditor){
if (!activeTextEditor) {
return undefined;
}
const position = activeTextEditor.selection.active;
Expand Down Expand Up @@ -177,3 +178,9 @@ export async function getCodeForRange(uri: vscode.Uri, range: vscode.Range): Pro

return undefined;
}

export function computeCodehash(fileContents: string): string {
// Normalize the file content by removing whitespace and newlines
const normalizedContent = fileContents.replace(/\s+/g, '');
return crypto.createHash('sha256').update(normalizedContent).digest('hex');
}
3 changes: 2 additions & 1 deletion vscode/src/tools/create/inline_code_completion.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ function replaceLineOfCode(line: number, replaceString: string) {
});
}

export async function createInlineCodeCompletion(geminiRepo: GeminiRepository) {
export async function createInlineCodeCompletion() {
// manual trigger using shortcut ctrl+space
logEvent('create-inline-code-completion', { 'type': "create" });
vscode.window.withProgress({
Expand Down Expand Up @@ -100,6 +100,7 @@ async function generateSuggestions(): Promise<string[]> {
const fileContent = editor.document.getText();
const filepath = editor.document.fileName; // don't include currentFile in most relevant files.
console.log("Current file path:", filepath);
///TODO: Replace with Generation Repository
var relevantFiles = await GeminiRepository.getInstance().findClosestDartFiles("Current file content:" + editor.document.getText() + "\n\n" + "Line of code:" + currentLineContent, undefined, true, filepath);
// Add code for all the elements used in the file.
const contextualCode = await new ContextualCodeProvider().getContextualCodeForCompletion(editor.document, getDartAnalyser());
Expand Down
77 changes: 76 additions & 1 deletion vscode/src/utilities/cache-manager.ts
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
import * as vscode from 'vscode';
import { logError } from './telemetry-reporter';

import { ContentEmbedding } from '@google/generative-ai';
import * as path from 'path';
import { computeCodehash } from '../shared/utils';
export class CacheManager {
private static instance: CacheManager;
private globalState: vscode.Memento;
Expand Down Expand Up @@ -58,4 +60,77 @@ export class CacheManager {
return 0;
}
}

async setGeminiCache(cacheData: { [filePath: string]: { codehash: string, embedding: ContentEmbedding } }): Promise<void> {
const excludePatterns = "**/{android,ios,web,linux,macos,windows,.dart_tool}/**";
const pubspecs = await vscode.workspace.findFiles("**/pubspec.yaml", excludePatterns);
if (pubspecs.length === 0) {
throw new Error("No pubspec.yaml found in the workspace.");
}
// Get all the flutter projects in the workspace
const flutterProjects = pubspecs.map((uri) => uri.fsPath);
const cache: { [flutterProject: string]: { [filePath: string]: { codehash: string, embedding: ContentEmbedding } } } = {};

// ITERATE OVER ALL THE FILES IN THE CACHE
// Find the flutter project for the file
// Add the file to the cache of that flutter project
for (const filePath in cacheData) {
const parentProjectPath = this.findParentFlutterProject(filePath, flutterProjects);
if (parentProjectPath) {
const key = "gemini-cache-" + computeCodehash(parentProjectPath);
await this.setGlobalValue(key, JSON.stringify(cacheData));
}
}

}

async getGeminiCache(): Promise<string | undefined> {
const excludePatterns = "**/{android,ios,web,linux,macos,windows,.dart_tool}/**";
const pubspecs = await vscode.workspace.findFiles("**/pubspec.yaml", excludePatterns);
if (pubspecs.length === 0) {
throw new Error("No pubspec.yaml found in the workspace.");
}

// Get all the flutter projects in the workspace
const flutterProjects = pubspecs.map((uri) => uri.fsPath);

const activeCache: { [filePath: string]: { codehash: string, embedding: ContentEmbedding } } = {};
// Return cache only for the parent flutter project of the current workspace
for (const projectPath of flutterProjects) {
const projectDir = path.dirname(projectPath);
const key = "gemini-cache-" + projectDir;
// await this.setGlobalValue<String>(key, "value");
const cacheString = await this.getGlobalValue<string>(key);
if (!cacheString) {
continue;
}
const cache = JSON.parse(cacheString);
// Add the cache for the project to the activeCache
Object.assign(activeCache, cache);
}

return JSON.stringify(activeCache);
}

// Helper function to find parent Flutter project
private findParentFlutterProject(filePath: string, flutterProjects: string[]) {
let parentProjectPath = null;
let maxCommonLength = -1;

for (const projectPath of flutterProjects) {
const projectDir = path.dirname(projectPath);

// Check if the current projectDir is a prefix of the filePath
if (filePath.startsWith(projectDir)) {
const commonLength = projectDir.length;

if (commonLength > maxCommonLength) {
maxCommonLength = commonLength;
parentProjectPath = projectDir;
}
}
}

return parentProjectPath;
}
}
3 changes: 2 additions & 1 deletion vscode/src/utilities/command-manager.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@ export function registerCommand(
}

export function initCommands(context: vscode.ExtensionContext, geminiRepo: any, analyzer: any, flutterGPTViewProvider: FlutterGPTViewProvider) {
const generationRepository: GenerationRepository = getUserPrefferedModel();
const generationRepository: GenerationRepository = geminiRepo; //TODO: Use up in the tree getUserPrefferedModel();
// List of commands to register, with their names and options.
const commands = [
{ name: 'dashai.attachToDash', handler: () => addToReference(context.globalState, flutterGPTViewProvider), options: { isCommand: true, isMenu: true, isShortcut: false } },
Expand All @@ -81,6 +81,7 @@ export function initCommands(context: vscode.ExtensionContext, geminiRepo: any,
{ name: 'dashai.optimizeCode', handler: (aiRepo: GenerationRepository, globalState: vscode.Memento, range: vscode.Range, anlyzer: ILspAnalyzer, elementName: string | undefined) => optimizeCode(generationRepository, context.globalState, range, anlyzer, elementName, context), options: { isCommand: true, isMenu: false, isShortcut: false } },
{ name: 'dashai.createInlineCodeCompletion', handler: () => createInlineCodeCompletion(geminiRepo), options: { isCommand: true, isMenu: true, isShortcut: true } },
{ name: 'dashai.clearChat', handler: () => flutterGPTViewProvider?.postMessageToWebview({ type: 'clearCommandDeck' }), options: { isCommand: true, isMenu: false, isShortcut: false } }

// Add more commands as needed.
];

Expand Down

0 comments on commit 6b01c47

Please sign in to comment.