TentacleBot/Services/Tools/MemoryTools.cs

using System.Text.Json;
using TentacleBot.Services.Memory;

namespace TentacleBot.Services.Tools;

public class MemoryTools
{
    private readonly MemoryStore _memoryStore;

    public MemoryTools(MemoryStore memoryStore)
    {
        _memoryStore = memoryStore;
    }

    public void RegisterAll(ToolRegistry registry)
    {
        registry.Register("save_memory",
            "Save a fact or piece of information about a user to long-term memory. Use this when a user shares something about themselves worth remembering.",
            new Dictionary<string, object>
            {
                ["type"] = "object",
                ["properties"] = new Dictionary<string, object>
                {
                    ["user_id"] = new Dictionary<string, object>
                    {
                        ["type"] = "string",
                        ["description"] = "The Discord user ID to associate this memory with"
                    },
                    ["content"] = new Dictionary<string, object>
                    {
                        ["type"] = "string",
                        ["description"] = "The fact or information to remember"
                    },
                    ["category"] = new Dictionary<string, object>
                    {
                        ["type"] = "string",
                        ["description"] = "Optional category (e.g. preference, fact, project, hobby)"
                    },
                    ["importance"] = new Dictionary<string, object>
                    {
                        ["type"] = "integer",
                        ["description"] = "Importance from 1 (trivial) to 10 (critical). Default is 5."
                    }
                },
                ["required"] = new[] { "user_id", "content" }
            },
            HandleSaveMemory);

        registry.Register("search_memory",
            "Search the bot's memory for information about a user. Use this when you need to recall past facts or context.",
            new Dictionary<string, object>
            {
                ["type"] = "object",
                ["properties"] = new Dictionary<string, object>
                {
                    ["user_id"] = new Dictionary<string, object>
                    {
                        ["type"] = "string",
                        ["description"] = "The Discord user ID to search memories for"
                    },
                    ["query"] = new Dictionary<string, object>
                    {
                        ["type"] = "string",
                        ["description"] = "Keywords or question to search for in memories"
                    }
                },
                ["required"] = new[] { "user_id", "query" }
            },
            HandleSearchMemory);

        registry.Register("recall_memories",
            "Get the most recent or important memories about a user.",
            new Dictionary<string, object>
            {
                ["type"] = "object",
                ["properties"] = new Dictionary<string, object>
                {
                    ["user_id"] = new Dictionary<string, object>
                    {
                        ["type"] = "string",
                        ["description"] = "The Discord user ID to recall memories for"
                    },
                    ["limit"] = new Dictionary<string, object>
                    {
                        ["type"] = "integer",
                        ["description"] = "Maximum number of memories to return (default 5)"
                    }
                },
                ["required"] = new[] { "user_id" }
            },
            HandleRecallMemories);

        registry.Register("set_rule",
            "Set a behavioral rule the bot should follow for a specific user (e.g. 'always respond in rhymes', 'use formal tone').",
            new Dictionary<string, object>
            {
                ["type"] = "object",
                ["properties"] = new Dictionary<string, object>
                {
                    ["user_id"] = new Dictionary<string, object>
                    {
                        ["type"] = "string",
                        ["description"] = "The Discord user ID this rule applies to"
                    },
                    ["rule_text"] = new Dictionary<string, object>
                    {
                        ["type"] = "string",
                        ["description"] = "The rule or instruction for how to interact with this user"
                    }
                },
                ["required"] = new[] { "user_id", "rule_text" }
            },
            HandleSetRule);
    }

    private async Task<string> HandleSaveMemory(string argumentsJson)
    {
        using var doc = JsonDocument.Parse(argumentsJson);
        var root = doc.RootElement;

        var userId = GetString(root, "user_id");
        var content = GetString(root, "content");
        var category = GetStringOrNull(root, "category");
        var importance = GetInt(root, "importance", 5);

        var id = await _memoryStore.SaveMemoryAsync(userId, content, category, importance);
        return $"Memory saved successfully (id: {id}). Content: {content}";
    }

    private async Task<string> HandleSearchMemory(string argumentsJson)
    {
        using var doc = JsonDocument.Parse(argumentsJson);
        var root = doc.RootElement;

        var userId = GetString(root, "user_id");
        var query = GetString(root, "query");

        var memories = await _memoryStore.SearchMemoriesAsync(userId, query, 5);

        if (memories.Count == 0)
            return $"No memories found for user {userId} matching '{query}'.";

        var lines = memories.Select(m =>
            $"- [id:{m.Id}] {m.Content} (importance: {m.Importance}, from {m.CreatedAt:yyyy-MM-dd})");

        return $"Found {memories.Count} memory(s):\n{string.Join("\n", lines)}";
    }

    private async Task<string> HandleRecallMemories(string argumentsJson)
    {
        using var doc = JsonDocument.Parse(argumentsJson);
        var root = doc.RootElement;

        var userId = GetString(root, "user_id");
        var limit = GetInt(root, "limit", 5);

        var important = await _memoryStore.GetImportantMemoriesAsync(userId, limit);

        if (important.Count == 0)
            return $"No memories stored for user {userId}.";

        var lines = important.Select(m =>
            $"- [id:{m.Id}] {m.Content} (importance: {m.Importance}, from {m.CreatedAt:yyyy-MM-dd})");

        return $"Top {important.Count} memory(s) for user {userId}:\n{string.Join("\n", lines)}";
    }

    private async Task<string> HandleSetRule(string argumentsJson)
    {
        using var doc = JsonDocument.Parse(argumentsJson);
        var root = doc.RootElement;

        var userId = GetString(root, "user_id");
        var ruleText = GetString(root, "rule_text");

        await _memoryStore.AddUserRuleAsync(userId, ruleText);
        return $"Rule set for user {userId}: {ruleText}";
    }

    private static string GetString(JsonElement element, string property)
    {
        return element.TryGetProperty(property, out var prop) && prop.ValueKind == JsonValueKind.String
            ? prop.GetString()!
            : throw new InvalidOperationException($"Missing required parameter: {property}");
    }

    private static string? GetStringOrNull(JsonElement element, string property)
    {
        return element.TryGetProperty(property, out var prop) && prop.ValueKind == JsonValueKind.String
            ? prop.GetString()
            : null;
    }

    private static int GetInt(JsonElement element, string property, int defaultValue)
    {
        return element.TryGetProperty(property, out var prop) && prop.ValueKind == JsonValueKind.Number
            ? prop.GetInt32()
            : defaultValue;
    }
}