memory

import { type PluginContext, tool } from "@lmstudio/sdk";
import { z } from "zod";
import { MemoryManager } from "./memoryManager";

const memoryManager = new MemoryManager();

// --- Tool Definitions ---

const readMemory = tool({
  name: "read_Memory",
  description: "Read memory slots. Call without arguments to list all slots, or specify key or id to read a specific slot.",
  parameters: {
    key: z.string().optional().describe("The key of the memory slot to read"),
    id: z.number().optional().describe("The ID of the memory slot to read"),
  },
  implementation: async ({ key, id }) => {
    return memoryManager.read(key, id);
  },
});

const addMemory = tool({
  name: "add_Memory",
  description: "Add a new memory slot with arbitrary data. Use this to remember information for later.",
  parameters: {
    key: z.string().describe("A unique key to identify this memory slot"),
    data: z.any().describe("The data to store (any JSON-serializable value)"),
  },
  implementation: async ({ key, data }) => {
    return memoryManager.add(key, data);
  },
});

const updateMemory = tool({
  name: "update_Memory",
  description: "Update an existing memory slot by ID. Can update key, data, or both.",
  parameters: {
    id: z.number().describe("The ID of the memory slot to update"),
    key: z.string().optional().describe("New key for the slot"),
    data: z.any().optional().describe("New data to store"),
  },
  implementation: async ({ id, key, data }) => {
    return memoryManager.update(id, key, data);
  },
});

const removeMemory = tool({
  name: "remove_Memory",
  description: "Permanently delete a memory slot by its ID.",
  parameters: {
    id: z.number().describe("The ID of the memory slot to delete"),
  },
  implementation: async ({ id }) => {
    return memoryManager.remove(id);
  },
});

const clearAllMemory = tool({
  name: "clear_All_Memory",
  description: "DANGER: Delete ALL memory slots. Use only if explicitly asked to wipe or reset memory.",
  parameters: {},
  implementation: async () => {
    return memoryManager.clear();
  },
});

// --- Main Entry Point ---

export async function main(context: PluginContext) {
  context.withToolsProvider(async () => {
    return [readMemory, addMemory, updateMemory, removeMemory, clearAllMemory];
  });
}