logging and caching

Author: pelikhan

packages/core/src/agent.ts

@@ -9,7 +9,33 @@ import { GenerationOptions } from "./generation"
 import { HTMLEscape } from "./html"
 import { prettifyMarkdown } from "./markdown"
 import { TraceOptions } from "./trace"
-import { logVerbose } from "./util"
+import { ellipse } from "./util"
+import debug from "debug"
+const dbg = debug("agent:memory")
+
+export type AgentMemoryCacheKey = { agent: string; query: string }
+export type AgentMemoryCacheValue = AgentMemoryCacheKey & {
+    answer: string
+    createdAt: number
+}
+export type AgentMemoryCache = WorkspaceFileCache<
+    AgentMemoryCacheKey,
+    AgentMemoryCacheValue
+>
+
+export function agentCreateCache(
+    options: Pick<GenerationOptions, "userState"> & { lookupOnly?: boolean }
+): AgentMemoryCache {
+    const cache = createCache<AgentMemoryCacheKey, AgentMemoryCacheValue>(
+        AGENT_MEMORY_CACHE_NAME,
+        {
+            type: "memory",
+            userState: options.userState,
+            lookupOnly: options.lookupOnly,
+        }
+    )
+    return cache
+}
 
 /**
  * Queries the agent's memory to retrieve contextual information relevant to a given query.
@@ -26,17 +52,19 @@ import { logVerbose } from "./util"
  * @returns Memory answer or undefined if no relevant memories are retrieved.
  */
 export async function agentQueryMemory(
+    cache: AgentMemoryCache,
     ctx: ChatGenerationContext,
     query: string,
-    options: Pick<GenerationOptions, "userState"> & Required<TraceOptions>
+    options: Required<TraceOptions>
 ) {
     if (!query) return undefined
 
-    const memories = await loadMemories(options)
+    const memories = await loadMemories(cache)
     if (!memories?.length) return undefined
 
     let memoryAnswer: string | undefined
     // always pre-query memory with cheap model
+    dbg(`query: ${query}`)
     const res = await ctx.runPrompt(
         async (_) => {
             _.$`Return the contextual information useful to answer <QUERY> from the content in <MEMORY>.
@@ -47,18 +75,21 @@ export async function agentQueryMemory(
                 "system"
             )
             _.def("QUERY", query)
-            await defMemory(_)
+            await defMemory(cache, _)
         },
         {
             model: "memory",
             system: [],
             flexTokens: AGENT_MEMORY_FLEX_TOKENS,
             label: "agent memory query",
+            cache: "agent_memory",
         }
     )
     if (!res.error)
         memoryAnswer = res.text.includes(TOKEN_NO_ANSWER) ? "" : res.text
-    else logVerbose(`agent memory query error: ${errorMessage(res.error)}`)
+    else dbg(`error: ${errorMessage(res.error)}`)
+
+    dbg(`answer: ${ellipse(memoryAnswer, 128)}`)
     return memoryAnswer
 }
 
@@ -73,25 +104,20 @@ export async function agentQueryMemory(
  * @param options - Configuration options, including user state and tracing details.
  */
 export async function agentAddMemory(
+    cache: AgentMemoryCache,
     agent: string,
     query: string,
     text: string,
-    options: Pick<GenerationOptions, "userState"> & Required<TraceOptions>
+    options: Required<TraceOptions>
 ) {
     const { trace } = options || {}
-    const cache = createCache<
-        { agent: string; query: string },
-        {
-            agent: string
-            query: string
-            answer: string
-        }
-    >(AGENT_MEMORY_CACHE_NAME)
-    const cacheKey = { agent, query }
-    const cachedValue = {
+    const cacheKey: AgentMemoryCacheKey = { agent, query }
+    const cachedValue: AgentMemoryCacheValue = {
         ...cacheKey,
         answer: text,
+        createdAt: Date.now(),
     }
+    dbg(`add ${agent}: ${ellipse(query, 80)} -> ${ellipse(text, 128)}`)
     await cache.set(cacheKey, cachedValue)
     trace.detailsFenced(
         `🧠 agent memory: ${HTMLEscape(query)}`,
@@ -100,19 +126,9 @@ export async function agentAddMemory(
     )
 }
 
-async function loadMemories(options: Pick<GenerationOptions, "userState">) {
-    const cache = createCache<
-        { agent: string; query: string },
-        {
-            agent: string
-            query: string
-            answer: string
-        }
-    >(AGENT_MEMORY_CACHE_NAME, {
-        lookupOnly: true,
-        userState: options.userState,
-    })
+async function loadMemories(cache: AgentMemoryCache) {
     const memories = await cache?.values()
+    memories?.sort((l, r) => l.createdAt - r.createdAt)
     return memories
 }
 
@@ -131,7 +147,11 @@ export async function traceAgentMemory(
     options: Pick<GenerationOptions, "userState"> & Required<TraceOptions>
 ) {
     const { trace } = options || {}
-    const memories = await loadMemories(options)
+    const cache = agentCreateCache({
+        userState: options.userState,
+        lookupOnly: true,
+    })
+    const memories = await loadMemories(cache)
     if (memories) {
         try {
             trace.startDetails("🧠 agent memory")
@@ -150,15 +170,10 @@ export async function traceAgentMemory(
     }
 }
 
-async function defMemory(ctx: ChatTurnGenerationContext) {
-    const cache = createCache<
-        { agent: string; query: string },
-        {
-            agent: string
-            query: string
-            answer: string
-        }
-    >(AGENT_MEMORY_CACHE_NAME)
+async function defMemory(
+    cache: AgentMemoryCache,
+    ctx: ChatTurnGenerationContext
+) {
     const memories = await cache.values()
     memories.reverse().forEach(({ agent, query, answer }, index) =>
         ctx.def(

packages/core/src/chatcache.ts

@@ -25,6 +25,7 @@ export function getChatCompletionCache(
     name?: string
 ): ChatCompletationRequestCache {
     return createCache<ChatCompletionRequestCacheKey, ChatCompletionResponse>(
-        name || CHAT_CACHE
+        name || CHAT_CACHE,
+        { type: "fs" }
     )
 }

packages/core/src/runpromptcontext.ts

@@ -87,7 +87,7 @@ import { resolveScript } from "./ast"
 import { dedent } from "./indent"
 import { runtimeHost } from "./host"
 import { writeFileEdits } from "./fileedits"
-import { agentAddMemory, agentQueryMemory } from "./agent"
+import { agentAddMemory, agentCreateCache, agentQueryMemory } from "./agent"
 import { YAMLStringify } from "./yaml"
 import { Project } from "./server/messages"
 import { mergeEnvVarsWithSystem, parametersToVars } from "./vars"
@@ -470,6 +470,7 @@ export function createChatGenerationContext(
         }
     }
 
+    const adbgm = debug(`agent:memory`)
     const defAgent = (
         name: string,
         description: string,
@@ -488,11 +489,12 @@ export function createChatGenerationContext(
             disableMemoryQuery,
             ...rest
         } = options || {}
-        const memory = !disableMemory
+        const memory = disableMemory
+            ? undefined
+            : agentCreateCache({ userState })
 
         name = name.replace(/^agent_/i, "")
         const adbg = debug(`agent:${name}`)
-        const adbgm = debug(`agent:${name}:memory`)
         adbg(`created ${variant || ""}`)
         const agentName = `agent_${name}${variant ? "_" + variant : ""}`
         const agentLabel = `agent ${name}${variant ? " " + variant : ""}`
@@ -547,12 +549,13 @@ export function createChatGenerationContext(
                 let memoryAnswer: string
                 if (memory && query && !disableMemoryQuery) {
                     memoryAnswer = await agentQueryMemory(
+                        memory,
                         ctx,
                         query +
                             (hasExtraArgs
                                 ? `\n${YAMLStringify(argsNoQuery)}`
                                 : ""),
-                        { userState, trace }
+                        { trace }
                     )
                     if (memoryAnswer) adbgm(`found ${memoryAnswer}`)
                 }
@@ -590,11 +593,11 @@ export function createChatGenerationContext(
                                 ) {
                                     adbgm(`add ${text}`)
                                     await agentAddMemory(
+                                        memory,
                                         agentName,
                                         query,
                                         text,
                                         {
-                                            userState,
                                             trace,
                                         }
                                     )
@@ -836,7 +839,8 @@ export function createChatGenerationContext(
                     ? TRANSCRIPTION_CACHE_NAME
                     : typeof cache === "string"
                       ? cache
-                      : undefined
+                      : undefined,
+                { type: "fs" }
             )
             if (cache) {
                 const hit = await _cache.getOrUpdate(

packages/core/src/vectorsearch.ts

@@ -56,7 +56,7 @@ export function createCachedEmbedder(
     const cache: EmbeddingsCache = createCache<
         EmbeddingsCacheKey,
         EmbeddingsResponse
-    >(cacheName || "embeddings")
+    >(cacheName || "embeddings", { type: "fs" })
 
     return async (inputs: string, cfg, options) => {
         const key: EmbeddingsCacheKey = {