added cache where we add promises and then we can check if those prom… (#5458)

* added cache where we add promises and then we can check if those promises are there and then we await them and we start them when the obj is made

* no console

* rm console

Author: chitalian

worker/src/lib/ai-gateway/AttemptBuilder.ts

@@ -123,7 +123,7 @@ export class AttemptBuilder {
 
     const attemptArrays = await Promise.all(
       providerData.map(async (data) => {
-        const byokAttempts = await this.buildByokAttempts(
+        const byokAttempts = this.buildByokAttempts(
           modelSpec,
           data,
           orgId,
@@ -132,13 +132,13 @@ export class AttemptBuilder {
         );
 
         // Always build PTB attempts (feature flag removed)
-        const ptbAttempts = await this.buildPtbAttempts(
+        const ptbAttempts = this.buildPtbAttempts(
           modelSpec,
           data,
           bodyMapping,
           plugins
         );
-        return [...byokAttempts, ...ptbAttempts];
+        return [...(await byokAttempts), ...(await ptbAttempts)];
       })
     );
 

worker/src/lib/ai-gateway/SimpleAIGateway.ts

@@ -81,7 +81,8 @@ export class SimpleAIGateway {
 
     const providerKeysManager = new ProviderKeysManager(
       new ProviderKeysStore(this.supabaseClient),
-      env
+      env,
+      this.orgId
     );
 
     // Create SecureCacheProvider for distributed caching

worker/src/lib/managers/ProviderKeysManager.ts

@@ -1,17 +1,26 @@
 import { ModelProviderName } from "@helicone-package/cost/models/providers";
 import { ProviderKey, ProviderKeysStore } from "../db/ProviderKeysStore";
-import {
-  getFromKVCacheOnly,
-  InMemoryCache,
-  removeFromCache,
-  storeInCache,
-} from "../util/cache/secureCache";
+import { getFromKVCacheOnly, storeInCache } from "../util/cache/secureCache";
 
 export class ProviderKeysManager {
+  providerKeysFromCache: Map<string, Promise<ProviderKey[] | null>> = new Map();
   constructor(
     private store: ProviderKeysStore,
-    private env: Env
-  ) {}
+    private env: Env,
+    orgId?: string
+  ) {
+    if (orgId) {
+      this.providerKeysFromCache.set(
+        orgId,
+        this.getProviderKeysFromCache(`provider_keys_${orgId}`)
+      );
+    }
+
+    this.providerKeysFromCache.set(
+      env.HELICONE_ORG_ID,
+      this.getProviderKeysFromCache(`provider_keys_${env.HELICONE_ORG_ID}`)
+    );
+  }
 
   async setProviderKeys() {
     const providerKeys = await this.store.getProviderKeys();
@@ -76,28 +85,34 @@ export class ProviderKeysManager {
     return filteredKeys[0];
   }
 
-  async getProviderKeys(orgId: string): Promise<ProviderKey[] | null> {
-    const kvCacheKey = `provider_keys_${orgId}`;
-    const memoryCacheKey = `provider_keys_in_memory_${orgId}`;
-
-    // Check in-memory cache first (fastest)
-    const cachedKeys =
-      InMemoryCache.getInstance<ProviderKey[]>().get(memoryCacheKey);
-    if (cachedKeys) {
-      return cachedKeys;
-    }
-
-    // Fall back to KV cache
+  private async getProviderKeysFromCache(
+    kvCacheKey: string
+  ): Promise<ProviderKey[] | null> {
     const kvKeys = await getFromKVCacheOnly(kvCacheKey, this.env, 43200);
     if (!kvKeys) {
       return null;
     }
+    return JSON.parse(kvKeys) as ProviderKey[];
+  }
+
+  async getProviderKeys(orgId: string): Promise<ProviderKey[] | null> {
+    if (this.providerKeysFromCache.has(orgId)) {
+      const promiseVal = this.providerKeysFromCache.get(orgId);
 
-    // Parse once and store in memory cache for subsequent requests
-    const parsedKeys = JSON.parse(kvKeys) as ProviderKey[];
-    InMemoryCache.getInstance<ProviderKey[]>().set(memoryCacheKey, parsedKeys);
+      if (promiseVal !== undefined) {
+        return await promiseVal;
+      }
+    }
 
-    return parsedKeys;
+    const result = await this.getProviderKeysFromCache(
+      `provider_keys_${orgId}`
+    );
+    if (result) {
+      this.providerKeysFromCache.set(orgId, Promise.resolve(result));
+      return result;
+    }
+
+    return null;
   }
 
   async getProviderKeyWithFetch(
@@ -107,6 +122,7 @@ export class ProviderKeysManager {
     keyCuid?: string
   ): Promise<ProviderKey | null> {
     const keys = await this.getProviderKeys(orgId);
+
     const validKey = this.chooseProviderKey(
       keys ?? [],
       provider,

worker/src/routers/api/apiRouter.ts

@@ -101,9 +101,11 @@ function getAPIRouterV1(
       }
 
       try {
-        const { removeFromCache } = await import("../../lib/util/cache/secureCache");
+        const { removeFromCache } = await import(
+          "../../lib/util/cache/secureCache"
+        );
         const cacheKeysToDelete: string[] = [];
-        
+
         if (data.versionId) {
           const promptBodyCacheKey = `prompt_body_${data.promptId}_${data.versionId}_${orgId}`;
           cacheKeysToDelete.push(promptBodyCacheKey);
@@ -118,30 +120,34 @@ function getAPIRouterV1(
         }
 
         await Promise.all(
-          cacheKeysToDelete.map(key => 
-            removeFromCache(key, env)
-          )
+          cacheKeysToDelete.map((key) => removeFromCache(key, env))
         );
 
-        return new Response(JSON.stringify({ 
-          success: true, 
-          deletedKeys: cacheKeysToDelete 
-        }), { 
-          status: 200,
-          headers: { "Content-Type": "application/json" }
-        });
+        return new Response(
+          JSON.stringify({
+            success: true,
+            deletedKeys: cacheKeysToDelete,
+          }),
+          {
+            status: 200,
+            headers: { "Content-Type": "application/json" },
+          }
+        );
       } catch (error) {
         console.error("Error resetting prompt cache:", error);
-        return new Response(JSON.stringify({ 
-          success: false, 
-          error: error instanceof Error ? error.message : "Unknown error" 
-        }), { 
-          status: 500,
-          headers: { "Content-Type": "application/json" }
-        });
+        return new Response(
+          JSON.stringify({
+            success: false,
+            error: error instanceof Error ? error.message : "Unknown error",
+          }),
+          {
+            status: 500,
+            headers: { "Content-Type": "application/json" },
+          }
+        );
       }
     }
-  )
+  );
 
   router.post(
     "/mock-set-provider-keys/:orgId",