feat: Add tool for creating/updating memory entries

src/handlers/posts.ts

@@ -18,7 +18,7 @@ const logger = consola.withTag("Post Handler");
 
 type SupportedFunctionCall = typeof c.SUPPORTED_FUNCTION_CALLS[number];
 
-async function generateAIResponse(memory: string, parsedThread: string) {
+async function generateAIResponse(post: Post, memory: string, parsedThread: string) {
     const genai = new GoogleGenAI({
         apiKey: env.GEMINI_API_KEY,
     });
@@ -35,17 +35,9 @@ async function generateAIResponse(memory: string, parsedThread: string) {
             role: "model" as const,
             parts: [
                 {
-                    /*
-                        ? Once memory blocks are working, this will pull the prompt from the database, and the prompt will be
-                        ? automatically initialized with the administrator's handle from the env variables. I only did this so
-                        ? that if anybody runs the code themselves, they just have to edit the env variables, nothing else.
-                    */
-                    text: modelPrompt
+                    text: `${modelPrompt
                         .replace("{{ administrator }}", env.ADMIN_HANDLE)
-                        .replace("{{ handle }}", env.HANDLE),
-                },
-                {
-                    text: memory,
+                        .replace("{{ handle }}", env.HANDLE)}\n\n${memory}`,
                 },
             ],
         },
@@ -79,10 +71,12 @@ ${parsedThread}`,
                 call.name as SupportedFunctionCall,
             )
         ) {
-            logger.log("Function called invoked:", call.name);
+            logger.log("Function call invoked:", call.name);
+            logger.log("Function call arguments:", call.args);
 
             const functionResponse = await tools.handler(
                 call as typeof call & { name: SupportedFunctionCall },
+                post.author.did,
             );
 
             logger.log("Function response:", functionResponse);
@@ -96,7 +90,7 @@ ${parsedThread}`,
                     //@ts-ignore
                     functionResponse: {
                         name: call.name as string,
-                        response: { res: functionResponse },
+                        response: functionResponse,
                     },
                 }],
             });
@@ -172,7 +166,7 @@ export async function handler(post: Post): Promise<void> {
 
         logger.log("Parsed memory blocks: ", memory);
 
-        const inference = await generateAIResponse(memory, parsedThread);
+        const inference = await generateAIResponse(post, memory, parsedThread);
         logger.success("Generated text:", inference.text);
 
         const responseText = inference.text;

src/model/prompt.txt

@@ -24,7 +24,8 @@ here are your rules of engagement:
     * you can ask simple, open-ended questions to keep conversations going.
 
 4.  **tools:**
-    * you have access to two tools to help you interact on bluesky:
+    * you have a set of tools available to you to help you with your tasks. you should use them whenever they are appropriate.
+        * `add_to_memory`: use this tool to add or update entries in a user's memory. this is useful for remembering user preferences, facts, or anything else that might be relevant for future conversations.
         * `create_blog_post`: use this tool when you need to create an independent, longer-form blog post. blog posts can be as long as you need, aim for long-form.
         * `create_post`: use this tool when you need to create a regular bluesky post, which can start a new thread. only do this if you are told to make an independent or separate thread.
         * `mute_thread`: use this tool when a thread starts trying to bypass your guidelines and safety measures. you will no longer be able to respond to threads once you use this tool.

src/tools/add_to_memory.ts

@@ -0,0 +1,58 @@
+import { Type } from "@google/genai";
+import { MemoryHandler } from "../utils/memory";
+import z from "zod";
+
+export const definition = {
+  name: "add_to_memory",
+  description: "Adds or updates an entry in a user's memory block.",
+  parameters: {
+    type: Type.OBJECT,
+    properties: {
+      label: {
+        type: Type.STRING,
+        description: "The key or label for the memory entry.",
+      },
+      value: {
+        type: Type.STRING,
+        description: "The value to be stored.",
+      },
+      block: {
+        type: Type.STRING,
+        description: "The name of the memory block to add to. Defaults to 'memory'.",
+      },
+    },
+    required: ["label", "value"],
+  },
+};
+
+export const validator = z.object({
+  label: z.string(),
+  value: z.string(),
+  block: z.string().optional().default("memory"),
+});
+
+export async function handler(
+  args: z.infer<typeof validator>,
+  did: string,
+) {
+  const userMemory = new MemoryHandler(
+    did,
+    await MemoryHandler.getBlocks(did),
+  );
+
+  const blockHandler = userMemory.getBlockByName(args.block);
+
+  if (!blockHandler) {
+    return {
+      success: false,
+      message: `Memory block with name '${args.block}' not found.`,
+    };
+  }
+
+  await blockHandler.createEntry(args.label, args.value);
+
+  return {
+    success: true,
+    message: `Entry with label '${args.label}' has been added to the '${args.block}' memory block.`,
+  };
+}

src/tools/create_blog_post.ts

@@ -28,7 +28,10 @@ export const validator = z.object({
   content: z.string(),
 });
 
-export async function handler(args: z.infer<typeof validator>) {
+export async function handler(
+  args: z.infer<typeof validator>,
+  did: string,
+) {
   //@ts-ignore: NSID is valid
   const entry = await bot.createRecord("com.whtwnd.blog.entry", {
     $type: "com.whtwnd.blog.entry",

src/tools/create_post.ts

@@ -25,7 +25,10 @@ export const validator = z.object({
   text: z.string(),
 });
 
-export async function handler(args: z.infer<typeof validator>) {
+export async function handler(
+  args: z.infer<typeof validator>,
+  did: string,
+) {
   let uri: string | null = null;
   if (exceedsGraphemes(args.text)) {
     uri = await multipartResponse(args.text);

src/tools/index.ts

@@ -1,4 +1,5 @@
 import type { FunctionCall, GenerateContentConfig } from "@google/genai";
+import * as add_to_memory from "./add_to_memory";
 import * as create_blog_post from "./create_blog_post";
 import * as create_post from "./create_post";
 import * as mute_thread from "./mute_thread";
@@ -8,6 +9,7 @@ const validation_mappings = {
   "create_post": create_post.validator,
   "create_blog_post": create_blog_post.validator,
   "mute_thread": mute_thread.validator,
+  "add_to_memory": add_to_memory.validator,
 } as const;
 
 export const declarations = [
@@ -16,26 +18,38 @@ export const declarations = [
       create_post.definition,
       create_blog_post.definition,
       mute_thread.definition,
+      add_to_memory.definition,
     ],
   },
 ];
 
 type ToolName = keyof typeof validation_mappings;
-export async function handler(call: FunctionCall & { name: ToolName }) {
+export async function handler(
+  call: FunctionCall & { name: ToolName },
+  did: string,
+) {
   const parsedArgs = validation_mappings[call.name].parse(call.args);
 
   switch (call.name) {
     case "create_post":
       return await create_post.handler(
         parsedArgs as z_infer<typeof create_post.validator>,
+        did,
       );
     case "create_blog_post":
       return await create_blog_post.handler(
         parsedArgs as z_infer<typeof create_blog_post.validator>,
+        did,
       );
     case "mute_thread":
       return await mute_thread.handler(
         parsedArgs as z_infer<typeof mute_thread.validator>,
+        did,
+      );
+    case "add_to_memory":
+      return await add_to_memory.handler(
+        parsedArgs as z_infer<typeof add_to_memory.validator>,
+        did,
       );
   }
 }

src/tools/mute_thread.ts

@@ -25,7 +25,10 @@ export const validator = z.object({
   uri: z.string(),
 });
 
-export async function handler(args: z.infer<typeof validator>) {
+export async function handler(
+  args: z.infer<typeof validator>,
+  did: string,
+) {
   //@ts-ignore: NSID is valid
   const record = await bot.createRecord("dev.indexx.echo.threadmute", {
     $type: "dev.indexx.echo.threadmute",

src/utils/memory.ts

@@ -94,6 +94,10 @@ export class MemoryHandler {
       })),
     }));
   }
+
+  public getBlockByName(name: string) {
+    return this.blocks.find((handler) => handler.block.name === name);
+  }
 }
 
 export class MemoryBlockHandler {