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

* fix(AI exports): add cahching layer to SQL exports * remove logs --------- Co-authored-by: Guy Ben-Aharon <baguy3@gmail.com>
Desion131525 · Nov 13, 2024 · e5dbbf2 · e5dbbf2
1 parent 959e540
commit e5dbbf2
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 2 deletions.
diff --git a/src/lib/data/export-metadata/export-sql-cache.ts b/src/lib/data/export-metadata/export-sql-cache.ts
@@ -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);
+};
diff --git a/src/lib/data/export-metadata/export-sql-script.ts b/src/lib/data/export-metadata/export-sql-script.ts
@@ -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;
 };
 

diff --git a/src/lib/utils.ts b/src/lib/utils.ts
@@ -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;
+};