fix(AI exports): add cahching layer to SQL exports (#390)

* fix(AI exports): add cahching layer to SQL exports

* remove logs

---------

Co-authored-by: Guy Ben-Aharon <baguy3@gmail.com>
This commit is contained in:
Jonathan Fishner
2024-11-13 17:18:27 +02:00
committed by GitHub
parent 959e5402b8
commit e5dbbf2eaa
3 changed files with 56 additions and 2 deletions

View File

@@ -0,0 +1,27 @@
import type { DatabaseType } from '@/lib/domain/database-type';
import { sha256 } from '@/lib/utils';
export const getFromCache = (key: string): string | null => {
try {
return localStorage.getItem(`sql-export-${key}`);
} catch (e) {
console.warn('Failed to read from localStorage:', e);
return null;
}
};
export const setInCache = (key: string, value: string): void => {
try {
localStorage.setItem(`sql-export-${key}`, value);
} catch (e) {
console.warn('Failed to write to localStorage:', e);
}
};
export const generateCacheKey = async (
databaseType: DatabaseType,
sqlScript: string
): Promise<string> => {
const rawKey = `${databaseType}:${sqlScript}`;
return await sha256(rawKey);
};

View File

@@ -3,6 +3,7 @@ import { OPENAI_API_KEY } from '@/lib/env';
import type { DatabaseType } from '@/lib/domain/database-type';
import type { DBTable } from '@/lib/domain/db-table';
import type { DataType } from '../data-types/data-types';
import { generateCacheKey, getFromCache, setInCache } from './export-sql-cache';
export const exportBaseSQL = (diagram: Diagram): string => {
const { tables, relationships } = diagram;
@@ -197,18 +198,27 @@ export const exportSQL = async (
signal?: AbortSignal;
}
): Promise<string> => {
const sqlScript = exportBaseSQL(diagram);
const cacheKey = await generateCacheKey(databaseType, sqlScript);
const cachedResult = getFromCache(cacheKey);
if (cachedResult) {
return cachedResult;
}
const [{ streamText, generateText }, { createOpenAI }] = await Promise.all([
import('ai'),
import('@ai-sdk/openai'),
]);
const openai = createOpenAI({
apiKey: OPENAI_API_KEY,
});
const sqlScript = exportBaseSQL(diagram);
const prompt = generateSQLPrompt(databaseType, sqlScript);
if (options?.stream) {
const { textStream, text } = await streamText({
const { textStream, text: textPromise } = await streamText({
model: openai('gpt-4o-mini-2024-07-18'),
prompt: prompt,
});
@@ -220,6 +230,9 @@ export const exportSQL = async (
options.onResultStream(textPart);
}
const text = await textPromise;
setInCache(cacheKey, text);
return text;
}
@@ -228,6 +241,7 @@ export const exportSQL = async (
prompt: prompt,
});
setInCache(cacheKey, text);
return text;
};

View File

@@ -180,3 +180,16 @@ export const cloneDiagram = (
updatedAt: new Date(),
};
};
export const sha256 = async (message: string): Promise<string> => {
const msgBuffer = new TextEncoder().encode(message);
const hashBuffer = await crypto.subtle.digest('SHA-256', msgBuffer);
const hashArray = Array.from(new Uint8Array(hashBuffer));
const hashHex = hashArray
.map((b) => b.toString(16).padStart(2, '0'))
.join('');
return hashHex;
};