Fixing duplicate message in cache artifacts, and cleaning up caching

logic

Author: jakedahn

jumble/integration/basic-flow.test.ts

@@ -22,7 +22,7 @@ import {
 const TOOLSHED_API_URL = Deno.env.get("TOOLSHED_API_URL") ??
   "http://localhost:8000/";
 const FRONTEND_URL = Deno.env.get("FRONTEND_URL") ?? "http://localhost:5173/";
-const HEADLESS = false;
+const HEADLESS = true;
 const ASTRAL_TIMEOUT = 60_000;
 
 console.log(`TOOLSHED_API_URL=${TOOLSHED_API_URL}`);

toolshed/routes/ai/llm/generateText.ts

@@ -23,14 +23,17 @@ export interface GenerateTextParams {
   abortSignal?: AbortSignal;
   max_tokens?: number;
   mode?: "json";
-  // Optional callback for when streaming is complete (used for caching)
-  onStreamComplete?: (
-    finalMessage: { role: "user" | "assistant"; content: string },
-  ) => void;
+  // Updated callback to receive complete data for caching
+  onStreamComplete?: (result: {
+    message: { role: "user" | "assistant"; content: string };
+    messages: { role: "user" | "assistant"; content: string }[];
+    originalRequest: GenerateTextParams;
+  }) => void;
 }
 
 export interface GenerateTextResult {
   message: { role: "user" | "assistant"; content: string };
+  messages: { role: "user" | "assistant"; content: string }[];
   stream?: ReadableStream;
 }
 
@@ -225,7 +228,10 @@ export async function generateText(
       messages[messages.length - 1].content = result;
     }
 
-    return { message: messages[messages.length - 1] };
+    return {
+      message: messages[messages.length - 1],
+      messages: [...messages],
+    };
   }
 
   // Create streaming response
@@ -271,10 +277,13 @@ export async function generateText(
         messages[messages.length - 1].content = result;
       }
 
-      // Call the onStreamComplete callback with the final message, if provided
-      // This is used to save the completed stream response to the cache
+      // Call the onStreamComplete callback with all the data needed for caching
       if (params.onStreamComplete) {
-        params.onStreamComplete(messages[messages.length - 1]);
+        params.onStreamComplete({
+          message: messages[messages.length - 1],
+          messages: [...messages],
+          originalRequest: params,
+        });
       }
 
       controller.close();
@@ -283,6 +292,7 @@ export async function generateText(
 
   return {
     message: messages[messages.length - 1],
+    messages: [...messages],
     stream,
   };
 }

toolshed/routes/ai/llm/llm.handlers.ts

@@ -92,17 +92,27 @@ export const getModels: AppRouteHandler<GetModelsRoute> = (c) => {
 export const generateText: AppRouteHandler<GenerateTextRoute> = async (c) => {
   const payload = await c.req.json();
 
-  console.log("LLM HASH", await cache.hashKey(JSON.stringify(payload)));
-  console.log("LLM PAYLOAD", JSON.stringify(payload, null, 2));
-
-  // Check cache for existing response
+  // First, check whether the request is cached, if so return the cached result
   const cacheKey = await cache.hashKey(JSON.stringify(payload));
   const cachedResult = await cache.loadItem(cacheKey);
   if (cachedResult) {
     const lastMessage = cachedResult.messages[cachedResult.messages.length - 1];
     return c.json(lastMessage);
   }
 
+  const persistCache = async (
+    messages: { role: string; content: string }[],
+  ) => {
+    try {
+      await cache.saveItem(cacheKey, {
+        ...payload,
+        messages,
+      });
+    } catch (e) {
+      console.error("Error saving response to cache:", e);
+    }
+  };
+
   const validationError = validateModelAndJsonMode(
     c,
     payload.model,
@@ -120,29 +130,17 @@ export const generateText: AppRouteHandler<GenerateTextRoute> = async (c) => {
       ...payload,
       abortSignal: c.req.raw.signal,
       max_tokens: payload.max_tokens || modelDefaultMaxTokens,
+      // If response is streaming, save to cache after the stream is complete
       onStreamComplete: payload.stream
-        ? (finalMessage) => {
-          // Save the completed stream response to the cache
-          cache.saveItem(cacheKey, {
-            ...payload,
-            messages: [...payload.messages, finalMessage],
-          }).catch((e) => {
-            console.error("Error saving streamed response to cache:", e);
-          });
+        ? async (result) => {
+          await persistCache(result.messages);
         }
         : undefined,
     });
 
+    // If response is not streaming, save to cache and return the message
     if (!payload.stream) {
-      // Save to cache
-      console.log("messagesssssss", [...payload.messages, result.message]);
-      console.log("=======================");
-      console.log("minus one", [...payload.messages]);
-      await cache.saveItem(cacheKey, {
-        ...payload,
-        // FIXME(jake): I believe this is persisting duplicate messages to the cached json blobs.
-        messages: [...payload.messages, result.message],
-      });
+      await persistCache(result.messages);
       return c.json(result.message);
     }