Implement memory management features with add, retrieve, and clear functionalities

Author: PrashamTrivedi

src/config/constants.ts

@@ -1,7 +1,9 @@
 import {join} from "jsr:@std/path"
+import Anthropic from "npm:@anthropic-ai/sdk"
+import {homedir} from "node:os"
 
-export const EDITOR_DIR = join(Deno.cwd(), "editor_dir")
-export const SESSIONS_DIR = join(Deno.cwd(), "sessions")
+export const EDITOR_DIR = join(homedir(), ".ComputerUseAgent", "editor_dir")
+export const SESSIONS_DIR = join(homedir(), ".ComputerUseAgent", "sessions")
 
 export const EDITOR_SYSTEM_PROMPT = Deno.env.get("EDITOR_SYSTEM_PROMPT") ??
     `You are a helpful assistant that helps users edit and work with text files.
@@ -22,4 +24,41 @@ export const API_CONFIG = {
     MODEL: "claude-3-5-sonnet-20241022",
     INTENT_MODEL: "claude-3-5-haiku-20241022",
     MAX_TOKENS: 4096,
+    MAX_INTENT_TOKENS: 20,
 }
+
+export const MEMORY_TOOLS: Anthropic.Beta.BetaToolUnion[] = [
+    {
+        name: "add_memory",
+        description: "Add a new memory to the system",
+        type: "custom",
+        input_schema: {
+            type: "object",
+            properties: {
+                content: {
+                    type: "string",
+                    description: "The content to store in memory",
+                },
+            },
+            required: ["content"],
+        },
+    },
+    {
+        type: "custom",
+        name: "get_memories",
+        description: "Retrieve all stored memories",
+        input_schema: {
+            type: "object",
+            properties: {},
+        },
+    },
+    {
+        type: "custom",
+        name: "clear_memories",
+        description: "Clear all stored memories",
+        input_schema: {
+            type: "object",
+            properties: {},
+        },
+    },
+]

src/main.ts

@@ -18,7 +18,7 @@ async function determineIntent(prompt: string): Promise<string> {
 
   const message = await client.messages.create({
     model: API_CONFIG.INTENT_MODEL,
-    max_tokens: 10,
+    max_tokens: API_CONFIG.MAX_INTENT_TOKENS,
     system: `You have following query from user, determine what is the intent of user. 
      and based on user's intent and available tools, pick the best tool that can help user.
      

src/modules/editor/editor_session.ts

@@ -1,15 +1,18 @@
 import {BaseSession} from "../../utils/session.ts"
-import {EDITOR_SYSTEM_PROMPT, API_CONFIG} from "../../config/constants.ts"
-import {ToolCall, ToolResult} from "../../types/interfaces.ts"
+import {EDITOR_SYSTEM_PROMPT, API_CONFIG, MEMORY_TOOLS} from "../../config/constants.ts"
+import {ToolResult} from "../../types/interfaces.ts"
 import {log} from "../../config/logging.ts"
 import {EditorHandlers} from "./handlers.ts"
+import {MemoryManager} from "../memory/memory_manager.ts"
 
 export class EditorSession extends BaseSession {
     private handlers: EditorHandlers
+    private memoryManager: MemoryManager
 
     constructor(sessionId?: string) {
         super(sessionId)
         this.handlers = new EditorHandlers()
+        this.memoryManager = new MemoryManager()
     }
 
     async processEdit(prompt: string): Promise<void> {
@@ -20,14 +23,17 @@ export class EditorSession extends BaseSession {
             }
             this.messages = [message]
 
-            log.info(`User input: ${JSON.stringify(message)}, System Prompt: ${EDITOR_SYSTEM_PROMPT}`)
+            log.info(`User input: ${JSON.stringify(message)}`)
 
             while (true) {
                 const response = await this.client.beta.messages.create({
                     model: API_CONFIG.MODEL,
                     max_tokens: API_CONFIG.MAX_TOKENS,
                     messages: this.messages,
-                    tools: [{type: "text_editor_20241022", name: "str_replace_editor"}],
+                    tools: [
+                        {type: "text_editor_20241022", name: "str_replace_editor"},
+                        ...MEMORY_TOOLS
+                    ],
                     system: EDITOR_SYSTEM_PROMPT,
                     betas: ["computer-use-2024-10-22"],
                 })
@@ -77,20 +83,35 @@ export class EditorSession extends BaseSession {
         const results: ToolResult[] = []
 
         for (const toolCall of content) {
-            if (toolCall.type === "tool_use" && toolCall.name === "str_replace_editor") {
-                log.info(`Editor tool call input: ${JSON.stringify(toolCall.input)}`)
-
-                const result = await this.handlers.handleTextEditorTool(toolCall.input)
-                const isError = "error" in result
-                const toolResultContent = isError
-                    ? [{type: "text", text: result.error}]
-                    : [{type: "text", text: result.content || ""}]
+            if (toolCall.type === "tool_use") {
+                let result
+                let isError = false
+
+                switch (toolCall.name) {
+                    case "str_replace_editor":
+                        result = await this.handlers.handleTextEditorTool(toolCall.input)
+                        isError = "error" in result
+                        break
+                    case "add_memory":
+                        result = await this.memoryManager.addMemory(toolCall.input.content)
+                        break
+                    case "get_memories":
+                        result = await this.memoryManager.getMemories()
+                        break
+                    case "clear_memories":
+                        await this.memoryManager.clearMemories()
+                        result = {message: "Memories cleared successfully"}
+                        break
+                }
 
                 results.push({
                     tool_call_id: toolCall.id,
                     output: {
                         type: "tool_result",
-                        content: toolResultContent,
+                        content: [{
+                            type: "text",
+                            text: JSON.stringify(result)
+                        }],
                         tool_use_id: toolCall.id,
                         is_error: isError,
                     },

src/modules/memory/memory_manager.ts

@@ -0,0 +1,41 @@
+import { Memory, MemoryFile } from "../../types/interfaces.ts";
+import { log } from "../../config/logging.ts";
+
+export class MemoryManager {
+    private memoryPath = "/root/memory.json";
+
+    private async readMemoryFile(): Promise<MemoryFile> {
+        try {
+            const content = await Deno.readTextFile(this.memoryPath);
+            return JSON.parse(content);
+        } catch {
+            return { memories: [] };
+        }
+    }
+
+    private async writeMemoryFile(data: MemoryFile): Promise<void> {
+        await Deno.writeTextFile(this.memoryPath, JSON.stringify(data, null, 2));
+    }
+
+    async addMemory(content: string): Promise<Memory> {
+        const memoryFile = await this.readMemoryFile();
+        const newMemory: Memory = {
+            id: crypto.randomUUID(),
+            content,
+            timestamp: Date.now(),
+        };
+        
+        memoryFile.memories.push(newMemory);
+        await this.writeMemoryFile(memoryFile);
+        return newMemory;
+    }
+
+    async getMemories(): Promise<Memory[]> {
+        const memoryFile = await this.readMemoryFile();
+        return memoryFile.memories;
+    }
+
+    async clearMemories(): Promise<void> {
+        await this.writeMemoryFile({ memories: [] });
+    }
+}

src/types/interfaces.ts

@@ -28,3 +28,13 @@ export interface SystemInfo {
   isWsl: boolean;
   shell: string;
 }
+
+export interface Memory {
+  id: string;
+  content: string;
+  timestamp: number;
+}
+
+export interface MemoryFile {
+  memories: Memory[];
+}