Ryanhub - file viewer
filename: tools/memory.go
branch: main
back to repo
package tools

import (
	"context"
	"encoding/json"
	"fmt"

	"assistant/memory"
)

var addMemoryParams = json.RawMessage(`{
  "type": "object",
  "properties": {
    "content": { "type": "string", "description": "Durable user fact or preference to remember." }
  },
  "required": ["content"],
  "additionalProperties": false
}`)

var listMemoriesParams = json.RawMessage(`{
  "type": "object",
  "properties": {
    "limit": { "type": "integer", "minimum": 1, "maximum": 200, "description": "Maximum memories to return (default 20)." }
  },
  "additionalProperties": false
}`)

var searchMemoriesParams = json.RawMessage(`{
  "type": "object",
  "properties": {
    "query": { "type": "string", "description": "Text to search for in stored memories." },
    "limit": { "type": "integer", "minimum": 1, "maximum": 200, "description": "Maximum memories to return (default 20)." }
  },
  "required": ["query"],
  "additionalProperties": false
}`)

var deleteMemoryParams = json.RawMessage(`{
  "type": "object",
  "properties": {
    "id": { "type": "integer", "description": "Memory id from list/search results." }
  },
  "required": ["id"],
  "additionalProperties": false
}`)

type addMemoryArgs struct {
	Content string `json:"content"`
}

type memoryListArgs struct {
	Limit int `json:"limit"`
}

type searchMemoryArgs struct {
	Query string `json:"query"`
	Limit int    `json:"limit"`
}

type deleteMemoryArgs struct {
	ID int64 `json:"id"`
}

func memoryTools(store *memory.Store) []Tool {
	return []Tool{
		{
			Name:        "add_memory",
			Description: "Store a durable long-term user memory for future conversations.",
			Parameters:  addMemoryParams,
			Run: func(ctx context.Context, args json.RawMessage) (string, error) {
				var a addMemoryArgs
				if err := json.Unmarshal(args, &a); err != nil {
					return "", err
				}
				id, err := store.AddMemory(a.Content)
				if err != nil {
					return "", err
				}
				return fmt.Sprintf("Saved memory #%d.", id), nil
			},
		},
		{
			Name:        "list_memories",
			Description: "List recent long-term memories.",
			Parameters:  listMemoriesParams,
			Run: func(ctx context.Context, args json.RawMessage) (string, error) {
				var a memoryListArgs
				if len(args) > 0 && string(args) != "null" {
					if err := json.Unmarshal(args, &a); err != nil {
						return "", err
					}
				}
				memories, err := store.ListMemories(a.Limit)
				if err != nil {
					return "", err
				}
				if len(memories) == 0 {
					return "No memories stored.", nil
				}
				b, err := json.MarshalIndent(memories, "", "  ")
				if err != nil {
					return "", err
				}
				return string(b), nil
			},
		},
		{
			Name:        "search_memories",
			Description: "Search long-term memories by text.",
			Parameters:  searchMemoriesParams,
			Run: func(ctx context.Context, args json.RawMessage) (string, error) {
				var a searchMemoryArgs
				if err := json.Unmarshal(args, &a); err != nil {
					return "", err
				}
				memories, err := store.SearchMemories(a.Query, a.Limit)
				if err != nil {
					return "", err
				}
				if len(memories) == 0 {
					return "No matching memories.", nil
				}
				b, err := json.MarshalIndent(memories, "", "  ")
				if err != nil {
					return "", err
				}
				return string(b), nil
			},
		},
		{
			Name:        "delete_memory",
			Description: "Delete one long-term memory by id.",
			Parameters:  deleteMemoryParams,
			Run: func(ctx context.Context, args json.RawMessage) (string, error) {
				var a deleteMemoryArgs
				if err := json.Unmarshal(args, &a); err != nil {
					return "", err
				}
				if err := store.DeleteMemory(a.ID); err != nil {
					return "", err
				}
				return fmt.Sprintf("Deleted memory %d.", a.ID), nil
			},
		},
	}
}