feat: add memory toolkit (#223)

This commit is contained in:
Max Novich
2024-11-07 15:42:05 -08:00
committed by GitHub
parent 808c84b8cd
commit be87308bc2
4 changed files with 422 additions and 1 deletions

View File

@@ -35,6 +35,7 @@ reasoner = "goose.toolkit.reasoner:Reasoner"
repo_context = "goose.toolkit.repo_context.repo_context:RepoContext"
synopsis = "goose.synopsis.toolkit:SynopsisDeveloper"
browser = "goose.toolkit.web_browser:BrowserToolkit"
memory = "goose.toolkit.memory:Memory"
[project.entry-points."goose.profile"]
default = "goose.profile:default_profile"
@@ -74,7 +75,8 @@ dev-dependencies = [
"mkdocstrings-python>=1.11.1",
"mkdocstrings>=0.26.1",
"pytest-mock>=3.14.0",
"pytest>=8.3.2"
"pytest>=8.3.2",
"tomli>=2.0.1",
]
[tool.uv.sources]

207
src/goose/toolkit/memory.py Normal file
View File

@@ -0,0 +1,207 @@
from pathlib import Path
from typing import Optional, List, Dict
import re
from jinja2 import Environment, FileSystemLoader
from goose.toolkit.base import Toolkit, tool
class Memory(Toolkit):
"""Memory toolkit for storing and retrieving natural
language memories with categories and tags"""
def __init__(self, *args: object, **kwargs: dict[str, object]) -> None:
super().__init__(*args, **kwargs)
# Setup memory directories
self.local_memory_dir = Path(".goose/memory")
self.global_memory_dir = Path.home() / ".config/goose/memory"
self._ensure_memory_dirs()
def _get_memories_data(self) -> dict:
"""Get memory data in a format suitable for template rendering"""
data = {"global": {}, "local": {}, "has_memories": False}
# Get global memories
if self.global_memory_dir.exists():
global_cats = [f.stem for f in self.global_memory_dir.glob("*.txt")]
for cat in sorted(global_cats):
memories = self._load_memories(cat, "global")
if memories:
data["global"][cat] = memories
data["has_memories"] = True
# Get local memories
if self.local_memory_dir.exists():
local_cats = [f.stem for f in self.local_memory_dir.glob("*.txt")]
for cat in sorted(local_cats):
memories = self._load_memories(cat, "local")
if memories:
data["local"][cat] = memories
data["has_memories"] = True
return data
def system(self) -> str:
"""Get the memory-specific additions to the system prompt"""
# Get the template directly since we need to render with our own variables
base_path = Path(__file__).parent / "prompts"
env = Environment(loader=FileSystemLoader(base_path))
template = env.get_template("memory.jinja")
return template.render(memories=self._get_memories_data())
def _ensure_memory_dirs(self) -> None:
"""Ensure memory directories exist"""
self.local_memory_dir.parent.mkdir(parents=True, exist_ok=True)
self.local_memory_dir.mkdir(exist_ok=True)
self.global_memory_dir.parent.mkdir(parents=True, exist_ok=True)
self.global_memory_dir.mkdir(exist_ok=True)
def _get_memory_file(self, category: str, scope: str = "global") -> Path:
"""Get the path to a memory category file"""
base_dir = self.global_memory_dir if scope == "global" else self.local_memory_dir
return base_dir / f"{category}.txt"
def _load_memories(self, category: str, scope: str = "global") -> List[Dict[str, str]]:
"""Load memories from a category file"""
memory_file = self._get_memory_file(category, scope)
if not memory_file.exists():
return []
memories = []
content = memory_file.read_text().strip()
if content:
for block in content.split("\n\n"):
if not block.strip():
continue
memory_lines = block.strip().split("\n")
tags = []
text = []
for line in memory_lines:
if line.startswith("#"):
tags.extend(tag.strip() for tag in line[1:].split())
else:
text.append(line)
memories.append({"text": "\n".join(text).strip(), "tags": tags})
return memories
def _save_memories(self, memories: List[Dict[str, str]], category: str, scope: str = "global") -> None:
"""Save memories to a category file"""
memory_file = self._get_memory_file(category, scope)
content = []
for memory in memories:
if memory["tags"]:
content.append(f"#{' '.join(memory['tags'])}")
content.append(memory["text"])
content.append("") # Empty line between memories
memory_file.write_text("\n".join(content))
@tool
def remember(self, text: str, category: str, tags: Optional[str] = None, scope: str = "global") -> str:
"""Save a memory with optional tags in a specific category
Args:
text (str): The memory text to store
category (str): The category to store the memory under (e.g., development, personal)
tags (str, optional): Space-separated tags to associate with the memory
scope (str): Where to store the memory - 'global' or 'local'
"""
# Clean and validate category name
category = re.sub(r"[^a-zA-Z0-9_-]", "_", category.lower())
# Process tags - remove any existing # prefix and store clean tags
tag_list = []
if tags:
tag_list = [tag.strip().lstrip("#") for tag in tags.split() if tag.strip()]
# Load existing memories
memories = self._load_memories(category, scope)
# Add new memory
memories.append({"text": text, "tags": tag_list})
# Save updated memories
self._save_memories(memories, category, scope)
tag_msg = f" with tags: {', '.join(tag_list)}" if tag_list else ""
return f"I'll remember that in the {category} category{tag_msg} ({scope} scope)"
@tool
def search(self, query: str, category: Optional[str] = None, scope: Optional[str] = None) -> str:
"""Search through memories by text and tags
Args:
query (str): Text to search for in memories and tags
category (str, optional): Specific category to search in
scope (str, optional): Which scope to search - 'global', 'local', or None (both)
"""
results = []
scopes = ["global", "local"] if scope is None else [scope]
for current_scope in scopes:
base_dir = self.global_memory_dir if current_scope == "global" else self.local_memory_dir
if not base_dir.exists():
continue
# Get categories to search
if category:
categories = [category]
else:
categories = [f.stem for f in base_dir.glob("*.txt")]
# Search in each category
for cat in categories:
memories = self._load_memories(cat, current_scope)
for memory in memories:
# Search in text and tags
if query.lower() in memory["text"].lower() or any(
query.lower() in tag.lower() for tag in memory["tags"]
):
tag_str = f" [tags: {', '.join(memory['tags'])}]" if memory["tags"] else ""
results.append(f"{current_scope}/{cat}: {memory['text']}{tag_str}")
if not results:
return "No matching memories found"
return "\n\n".join(results)
@tool
def list_categories(self, scope: Optional[str] = None) -> str:
"""List all memory categories
Args:
scope (str, optional): Which scope to list - 'global', 'local', or None (both)
"""
categories = []
if scope in (None, "local") and self.local_memory_dir.exists():
local_cats = [f.stem for f in self.local_memory_dir.glob("*.txt")]
if local_cats:
categories.append("Local categories:")
categories.extend(f" - {cat}" for cat in sorted(local_cats))
if scope in (None, "global") and self.global_memory_dir.exists():
global_cats = [f.stem for f in self.global_memory_dir.glob("*.txt")]
if global_cats:
categories.append("Global categories:")
categories.extend(f" - {cat}" for cat in sorted(global_cats))
if not categories:
return "No categories found in the specified scope(s)"
return "\n".join(categories)
@tool
def forget_category(self, category: str, scope: str = "global") -> str:
"""Remove an entire category of memories
Args:
category (str): The category to remove
scope (str): Which scope to remove from - 'global' or 'local'
"""
memory_file = self._get_memory_file(category, scope)
if not memory_file.exists():
return f"No {category} category found in {scope} scope"
memory_file.unlink()
return f"Successfully removed {category} category from {scope} scope"

View File

@@ -0,0 +1,69 @@
I have access to a memory system that helps me store and recall information across conversations. The memories are organized into categories and can be tagged for better searchability.
I can:
1. Remember information in categories (like "development", "preferences", "personal") with optional tags
2. Search through memories by text or tags
3. List all memory categories
4. Remove entire categories of memories
When users share important information like:
- Their name or personal details
- Project preferences
- Common tasks or workflows
- Configuration settings
I should:
1. Identify the key piece of information
2. Ask if they'd like me to remember it for future reference
3. If they agree:
- Suggest an appropriate category (e.g., "personal" for preferences, "development" for coding practices)
- Ask if they want any specific tags for easier searching
- Ask whether to store it:
- Locally (.goose/memory) for project-specific information
- Globally (~/.config/goose/memory) for user-wide preferences
4. Use the remember tool with the chosen category, tags, and scope
Example:
User: "For this project, we use black for code formatting"
Assistant: "I notice you mentioned a development preference. Would you like me to remember this for future conversations?"
User: "Yes please"
Assistant: "I'll store this in the 'development' category. Would you like me to add any specific tags? For example: #formatting #tools"
User: "Yes, those tags work"
Assistant: "And should I store this locally for just this project, or globally for all projects?"
User: "Locally please"
Assistant: *uses remember tool with category="development", tags="formatting tools", scope="local"*
{% if memories.has_memories %}
Here are the existing memories I have access to:
{% if memories.global %}
Global memories:
{% for category, category_memories in memories.global.items() %}
Category: {{ category }}
{% for memory in category_memories %}
- {{ memory.text }}{% if memory.tags %} [tags: {{ memory.tags|join(' ') }}]{% endif %}
{% endfor %}
{% endfor %}
{% endif %}
{% if memories.local %}
Local memories:
{% for category, category_memories in memories.local.items() %}
Category: {{ category }}
{% for memory in category_memories %}
- {{ memory.text }}{% if memory.tags %} [tags: {{ memory.tags|join(' ') }}]{% endif %}
{% endfor %}
{% endfor %}
{% endif %}
{% else %}
No existing memories found.
{% endif %}
I should always:
- Ask before storing information
- Suggest relevant categories and tags
- Clarify the storage scope
- Provide feedback about what was stored

View File

@@ -0,0 +1,143 @@
from unittest.mock import MagicMock
import pytest
from goose.toolkit.memory import Memory
@pytest.fixture
def memory_toolkit(tmp_path):
"""Create a memory toolkit instance with temporary directories"""
mock_notifier = MagicMock()
toolkit = Memory(notifier=mock_notifier)
# Override memory directories for testing
toolkit.local_memory_dir = tmp_path / ".goose/memory"
toolkit.global_memory_dir = tmp_path / ".config/goose/memory"
toolkit._ensure_memory_dirs()
return toolkit
def test_remember_global(memory_toolkit):
"""Test storing a memory in global scope"""
result = memory_toolkit.remember("Test memory", "test_category", tags="tag1 tag2", scope="global")
assert "test_category" in result
assert "tag1" in result
assert "tag2" in result
assert "global" in result
# Verify file content
memory_file = memory_toolkit.global_memory_dir / "test_category.txt"
content = memory_file.read_text()
assert "#tag1 tag2" in content
assert "Test memory" in content
def test_remember_local(memory_toolkit):
"""Test storing a memory in local scope"""
result = memory_toolkit.remember("Local test", "local_category", tags="local", scope="local")
assert "local_category" in result
assert "local" in result
# Verify file content
memory_file = memory_toolkit.local_memory_dir / "local_category.txt"
content = memory_file.read_text()
assert "#local" in content
assert "Local test" in content
def test_search_by_text(memory_toolkit):
"""Test searching memories by text"""
memory_toolkit.remember("Test memory one", "category1", scope="global")
memory_toolkit.remember("Test memory two", "category2", scope="global")
result = memory_toolkit.search("memory")
assert "Test memory one" in result
assert "Test memory two" in result
def test_search_by_tag(memory_toolkit):
"""Test searching memories by tag"""
memory_toolkit.remember("Tagged memory", "tagged", tags="findme test", scope="global")
memory_toolkit.remember("Another tagged", "tagged", tags="findme other", scope="global")
result = memory_toolkit.search("findme")
assert "Tagged memory" in result
assert "Another tagged" in result
def test_search_specific_category(memory_toolkit):
"""Test searching in a specific category"""
memory_toolkit.remember("Memory in cat1", "cat1", scope="global")
memory_toolkit.remember("Memory in cat2", "cat2", scope="global")
result = memory_toolkit.search("Memory", category="cat1")
assert "Memory in cat1" in result
assert "Memory in cat2" not in result
def test_list_categories(memory_toolkit):
"""Test listing memory categories"""
memory_toolkit.remember("Global memory", "global_cat", scope="global")
memory_toolkit.remember("Local memory", "local_cat", scope="local")
result = memory_toolkit.list_categories()
assert "global_cat" in result
assert "local_cat" in result
# Test scope filtering
global_only = memory_toolkit.list_categories(scope="global")
assert "global_cat" in global_only
assert "local_cat" not in global_only
def test_forget_category(memory_toolkit):
"""Test removing a category"""
memory_toolkit.remember("Memory to forget", "forget_me", scope="global")
assert (memory_toolkit.global_memory_dir / "forget_me.txt").exists()
result = memory_toolkit.forget_category("forget_me", scope="global")
assert "Successfully removed" in result
assert not (memory_toolkit.global_memory_dir / "forget_me.txt").exists()
def test_invalid_category_name(memory_toolkit):
"""Test that invalid category names are sanitized"""
result = memory_toolkit.remember("Test memory", "test/category!", tags="tag", scope="global")
assert "test_category_" in result
# Verify file was created with sanitized name
files = list(memory_toolkit.global_memory_dir.glob("*.txt"))
assert len(files) == 1
assert "test_category_" in files[0].name
def test_system_prompt_includes_memories(memory_toolkit):
"""Test that the system prompt includes existing memories"""
# Add some test memories
memory_toolkit.remember("Global test memory", "global_cat", tags="tag1 tag2", scope="global")
memory_toolkit.remember("Local test memory", "local_cat", tags="tag3", scope="local")
system_prompt = memory_toolkit.system()
# Check that the base prompt is included
assert "I have access to a memory system" in system_prompt
# Check that memories are included
assert "Global memories:" in system_prompt
assert "Category: global_cat" in system_prompt
assert "Global test memory" in system_prompt
assert "[tags: tag1 tag2]" in system_prompt
assert "Local memories:" in system_prompt
assert "Category: local_cat" in system_prompt
assert "Local test memory" in system_prompt
assert "[tags: tag3]" in system_prompt
def test_system_prompt_empty_memories(memory_toolkit):
"""Test that the system prompt handles no existing memories gracefully"""
system_prompt = memory_toolkit.system()
# Check that the base prompt is included
assert "I have access to a memory system" in system_prompt
# Check that empty memory state is handled
assert "No existing memories found" in system_prompt