feat: [ENG-3265] cache layer for prompt bodies on AI Gateway (#5009)

connortbot · web-flow · commit 9fad9c6b21c8 · 2025-10-14T13:21:56.000-07:00
* cache prompt bodies

* warning

* instant cache clearing on environment changes and deletes

* remove redundant cache reset

* production impl for prompt body caching

* remove console logs

* remove console log
diff --git a/valhalla/jawn/src/lib/clients/cloudflareKV.ts b/valhalla/jawn/src/lib/clients/cloudflareKV.ts
@@ -132,6 +132,37 @@ export async function removeFromCache(key: string): Promise<void> {
   );
 }
 
+export async function removeSecureCacheEntries(
+  keys: string[]
+): Promise<void> {
+  const namespaceId = process.env.CLOUDFLARE_KV_NAMESPACE_ID ?? "";
+  const accountId = process.env.CLOUDFLARE_ACCOUNT_ID ?? "";
+
+  // since on worker, we have secure cache (both HMAC key1 and HMAC key2)
+  await Promise.all(
+    keys
+      .filter((key) => Boolean(key))
+      .map(async (key) => {
+        const hashedKeys = await Promise.all([
+          hashWithHmac(key, 1),
+          hashWithHmac(key, 2),
+        ]);
+
+        await Promise.all(
+          hashedKeys.map((hashedKey) =>
+            cloudflare.kv.namespaces.values.delete(
+              namespaceId,
+              hashedKey,
+              {
+                account_id: accountId,
+              }
+            )
+          )
+        );
+      })
+  );
+}
+
 export async function storeInCache(
   key: string,
   value: string,
diff --git a/valhalla/jawn/src/lib/resetPromptCache.ts b/valhalla/jawn/src/lib/resetPromptCache.ts
@@ -0,0 +1,105 @@
+import { MAX_RETRIES } from "./refetchKeys";
+import { removeSecureCacheEntries } from "./clients/cloudflareKV";
+import { ENVIRONMENT } from "./clients/constant";
+
+type ResetPromptCacheParams = {
+  orgId: string;
+  promptId: string;
+  versionId?: string;
+  environment?: string;
+};
+
+const buildPromptCacheKeys = ({
+  orgId,
+  promptId,
+  versionId,
+  environment,
+}: ResetPromptCacheParams): string[] => {
+  const keys = new Set<string>();
+
+  // Default production scope
+  if (!versionId && !environment) {
+    keys.add(`prompt_version_${promptId}_prod_${orgId}`);
+  }
+
+  if (versionId) {
+    keys.add(`prompt_version_${promptId}_version:${versionId}_${orgId}`);
+    keys.add(`prompt_body_${promptId}_${versionId}_${orgId}`);
+    keys.add(`prompt_version_${promptId}_prod_${orgId}`);
+  }
+
+  if (environment) {
+    keys.add(`prompt_version_${promptId}_env:${environment}_${orgId}`);
+    if (environment === "production") {
+      keys.add(`prompt_version_${promptId}_prod_${orgId}`);
+    }
+  }
+
+  return Array.from(keys);
+};
+
+async function resetPromptCacheDev(
+  { orgId, promptId, versionId, environment }: ResetPromptCacheParams,
+  retries = MAX_RETRIES
+) {
+  try {
+    const res = await fetch(
+      `${process.env.HELICONE_WORKER_API}/reset-prompt-cache/${orgId}`,
+      {
+        method: "POST",
+        headers: {
+          "Content-Type": "application/json",
+        },
+        body: JSON.stringify({
+          promptId,
+          versionId,
+          environment,
+        }),
+      }
+    );
+
+    if (!res.ok) {
+      console.error(res);
+      if (retries > 0) {
+        await new Promise((resolve) =>
+          setTimeout(resolve, 10_000 * (MAX_RETRIES - retries))
+        );
+        await resetPromptCacheDev(
+          { orgId, promptId, versionId, environment },
+          retries - 1
+        );
+      }
+    }
+  } catch (error) {
+    console.error(error);
+    if (retries > 0) {
+      await new Promise((resolve) =>
+        setTimeout(resolve, 10_000 * (MAX_RETRIES - retries))
+      );
+      await resetPromptCacheDev(
+        { orgId, promptId, versionId, environment },
+        retries - 1
+      );
+    }
+  }
+}
+
+export async function resetPromptCache(
+  params: ResetPromptCacheParams
+): Promise<void> {
+  if (!params.promptId || !params.orgId) {
+    console.warn("Missing promptId or orgId when resetting prompt cache.");
+    return;
+  }
+
+  if (ENVIRONMENT === "production") {
+    const cacheKeys = buildPromptCacheKeys(params);
+    if (cacheKeys.length === 0) {
+      return;
+    }
+    await removeSecureCacheEntries(cacheKeys);
+    return;
+  }
+
+  await resetPromptCacheDev(params);
+}
diff --git a/valhalla/jawn/src/managers/prompt/PromptManager.ts b/valhalla/jawn/src/managers/prompt/PromptManager.ts
@@ -28,8 +28,8 @@ import { RequestManager } from "../request/RequestManager";
 import { S3Client } from "../../lib/shared/db/s3Client";
 import type { OpenAIChatRequest } from "@helicone-package/llm-mapper/mappers/openai/chat-v2";
 import { AuthParams } from "../../packages/common/auth/types";
-import { StringChain } from "lodash";
 import { Prompt2025Input } from "../../lib/db/ClickhouseWrapper";
+import { resetPromptCache as invalidatePromptCache } from "../../lib/resetPromptCache";
 
 
 const PROMPT_ID_LENGTH = 6;
@@ -50,6 +50,21 @@ export class Prompt2025Manager extends BaseManager {
     );
   }
 
+  private async resetPromptCache(params: {
+    promptId: string;
+    versionId?: string;
+    environment?: string;
+  }): Promise<void> {
+    try {
+      await invalidatePromptCache({
+        orgId: this.authParams.organizationId,
+        ...params,
+      });
+    } catch (error) {
+      console.error("Error resetting prompt cache:", error);
+    }
+  }
+
   private generateRandomPromptId() : string {
     const chars = '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ';
     let result = '';
@@ -666,6 +681,12 @@ export class Prompt2025Manager extends BaseManager {
     if (updateEnvResult.error) {
       return err(updateEnvResult.error);
     }
+
+    await this.resetPromptCache({
+      promptId: params.promptId,
+      environment: params.environment,
+    });
+
     return ok(null);
   }
 
@@ -709,6 +730,11 @@ export class Prompt2025Manager extends BaseManager {
       return err(result.error);
     }
 
+    // remove prod cache
+    await this.resetPromptCache({
+      promptId: params.promptId,
+    });
+
     return ok(null);
   }
 
@@ -730,6 +756,11 @@ export class Prompt2025Manager extends BaseManager {
       return err(s3Result.error);
     }
 
+    await this.resetPromptCache({
+      promptId: params.promptId,
+      versionId: params.promptVersionId
+    });
+
     return ok(null);
   }
 
diff --git a/worker/src/lib/managers/PromptManager.ts b/worker/src/lib/managers/PromptManager.ts
@@ -15,6 +15,51 @@ export class PromptManager {
     private env: Env
   ) {}
 
+  private buildPromptVersionCacheKey(
+    params: HeliconeChatCreateParams,
+    orgId: string
+  ): string | null {
+    if (!params.prompt_id) {
+      return null;
+    }
+
+    const scope = params.environment
+      ? `env:${params.environment}`
+      : params.version_id
+      ? `version:${params.version_id}`
+      : "prod";
+
+    return `prompt_version_${params.prompt_id}_${scope}_${orgId}`;
+  }
+
+  private buildPromptBodyCacheKey(
+    promptId: string,
+    versionId: string,
+    orgId: string
+  ): string {
+    return `prompt_body_${promptId}_${versionId}_${orgId}`;
+  }
+
+  private async getPromptVersionIdWithCache(
+    params: HeliconeChatCreateParams,
+    orgId: string
+  ): Promise<Result<string, string>> {
+    const cacheKey = this.buildPromptVersionCacheKey(params, orgId);
+    if (!cacheKey) {
+      return this.promptStore.getPromptVersionId(params, orgId);
+    }
+
+    return await getAndStoreInCache(
+      cacheKey,
+      this.env,
+      async () => {
+        return this.promptStore.getPromptVersionId(params, orgId);
+      },
+      300,
+      false
+    );
+  }
+
   async getSourcePromptBodyWithFetch(
     params: HeliconeChatCreateParams,
     orgId: string
@@ -27,14 +72,23 @@ export class PromptManager {
       string
     >
   > {
-    const versionIdResult = await this.promptStore.getPromptVersionId(
+    const versionIdResult = await this.getPromptVersionIdWithCache(
       params,
       orgId
     );
     if (isErr(versionIdResult)) return err(versionIdResult.error);
 
+    const promptId = params.prompt_id;
+    if (!promptId) {
+      return err("No prompt ID provided");
+    }
+
     return await getAndStoreInCache(
-      `prompt_body_${versionIdResult.data}_${orgId}`,
+      this.buildPromptBodyCacheKey(
+        promptId,
+        versionIdResult.data,
+        orgId
+      ),
       this.env,
       async () => {
         try {
@@ -50,7 +104,7 @@ export class PromptManager {
           return err(`Error retrieving prompt body: ${error}`);
         }
       },
-      undefined,
+      300,
       false
     );
   }
diff --git a/worker/src/routers/api/apiRouter.ts b/worker/src/routers/api/apiRouter.ts
@@ -78,6 +78,71 @@ function getAPIRouterV1(
     }
   );
 
+  router.post(
+    "/reset-prompt-cache/:orgId",
+    async (
+      { params: { orgId } },
+      requestWrapper: RequestWrapper,
+      env: Env,
+      ctx: ExecutionContext
+    ) => {
+      if (env.ENVIRONMENT !== "development") {
+        return new Response("not allowed", { status: 403 });
+      }
+
+      const data = await requestWrapper.unsafeGetJson<{
+        promptId: string;
+        versionId?: string;
+        environment?: string;
+      }>();
+
+      if (!data || !data.promptId) {
+        return new Response("promptId is required", { status: 400 });
+      }
+
+      try {
+        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);
+
+          const promptVersionCacheKey = `prompt_version_${data.promptId}_version:${data.versionId}_${orgId}`;
+          cacheKeysToDelete.push(promptVersionCacheKey);
+        }
+
+        if (data.environment) {
+          const promptVersionCacheKey = `prompt_version_${data.promptId}_env:${data.environment}_${orgId}`;
+          cacheKeysToDelete.push(promptVersionCacheKey);
+        }
+
+        await Promise.all(
+          cacheKeysToDelete.map(key => 
+            removeFromCache(key, env)
+          )
+        );
+
+        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" }
+        });
+      }
+    }
+  )
+
   router.post(
     "/mock-set-provider-keys/:orgId",
     async (
