484 lines
18 KiB
Python
484 lines
18 KiB
Python
"""Memory Store - Hierarchical file-based memory storage with MCP support."""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import re
|
|
import sys
|
|
import hashlib
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Optional, List
|
|
|
|
# === Hierarchical Memory Store ===
|
|
|
|
class MemoryStore:
|
|
"""Hierarchical file-based memory storage with categories."""
|
|
|
|
def __init__(self, memory_dir: Path):
|
|
self.memory_dir = Path(memory_dir)
|
|
self.memory_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
def _generate_id(self, title: str) -> str:
|
|
return hashlib.md5(title.encode()).hexdigest()[:8]
|
|
|
|
def _parse_frontmatter(self, content: str) -> tuple[dict, str]:
|
|
pattern = r'^---\n(.*?)\n---\n?(.*)$'
|
|
match = re.match(pattern, content, re.DOTALL)
|
|
if match:
|
|
fm_text = match.group(1)
|
|
body = match.group(2)
|
|
fm = {}
|
|
for line in fm_text.split('\n'):
|
|
if ':' in line:
|
|
key, value = line.split(':', 1)
|
|
key = key.strip()
|
|
value = value.strip().strip('[]"\'')
|
|
if value.startswith('['):
|
|
fm[key] = [v.strip().strip('"\'') for v in value[1:-1].split(',') if v.strip()]
|
|
else:
|
|
fm[key] = value
|
|
return fm, body
|
|
return {}, content
|
|
|
|
def _create_frontmatter(self, data: dict, body: str) -> str:
|
|
lines = ["---"]
|
|
for key, value in sorted(data.items()):
|
|
if isinstance(value, list):
|
|
lines.append(f"{key}: [{', '.join(value)}]")
|
|
else:
|
|
lines.append(f"{key}: {value}")
|
|
lines.append("---\n")
|
|
if body.strip():
|
|
lines.append(body.strip())
|
|
return "\n".join(lines)
|
|
|
|
def _get_category_path(self, category: str = None) -> Path:
|
|
if not category or category == "/":
|
|
return self.memory_dir
|
|
path = self.memory_dir / category.lstrip("/")
|
|
path.mkdir(parents=True, exist_ok=True)
|
|
return path
|
|
|
|
def _get_memory_path(self, category: str, memory_id: str) -> Path:
|
|
return self._get_category_path(category) / f"{memory_id}.md"
|
|
|
|
def _relpath(self, path: Path) -> str:
|
|
"""Get relative path from memory_dir."""
|
|
try:
|
|
return str(path.relative_to(self.memory_dir)).replace("\\", "/")
|
|
except ValueError:
|
|
return str(path)
|
|
|
|
def add(self, title: str, content: str, category: str = None, tags: List[str] = None, id: str = None) -> str:
|
|
"""Add memory to category."""
|
|
memory_id = id or self._generate_id(title)
|
|
filepath = self._get_memory_path(category, memory_id)
|
|
|
|
now = datetime.now().isoformat()
|
|
data = {
|
|
"id": memory_id,
|
|
"title": title,
|
|
"created": now,
|
|
"updated": now,
|
|
"tags": tags or [],
|
|
"category": category or "/"
|
|
}
|
|
|
|
filepath.write_text(self._create_frontmatter(data, content), encoding="utf-8")
|
|
return memory_id
|
|
|
|
def get(self, memory_id: str, category: str = None) -> Optional[dict]:
|
|
"""Get memory by ID, optionally in category."""
|
|
if category:
|
|
filepath = self._get_memory_path(category, memory_id)
|
|
if filepath.exists():
|
|
return self._read_memory(filepath)
|
|
return None
|
|
|
|
# Search all categories
|
|
for md_file in self.memory_dir.rglob("*.md"):
|
|
fm, _ = self._parse_frontmatter(md_file.read_text(encoding="utf-8"))
|
|
if fm.get("id") == memory_id:
|
|
result = fm
|
|
result["content"] = self._read_memory(md_file)["content"]
|
|
return result
|
|
return None
|
|
|
|
def _read_memory(self, filepath: Path) -> dict:
|
|
fm, body = self._parse_frontmatter(filepath.read_text(encoding="utf-8"))
|
|
fm["content"] = body.strip()
|
|
fm["category"] = self._relpath(filepath.parent)
|
|
return fm
|
|
|
|
def list(self, category: str = None, recursive: bool = True) -> List[dict]:
|
|
"""List memories in category."""
|
|
memories = []
|
|
cat_path = self._get_category_path(category)
|
|
|
|
pattern = "*.md" if recursive else "*.md"
|
|
for md_file in cat_path.glob(pattern):
|
|
if md_file.name.startswith("_"):
|
|
continue
|
|
fm, _ = self._parse_frontmatter(md_file.read_text(encoding="utf-8"))
|
|
fm["category"] = self._relpath(md_file.parent)
|
|
memories.append(fm)
|
|
|
|
return sorted(memories, key=lambda x: (x.get('category', ''), x.get('updated', '')), reverse=True)
|
|
|
|
def search(self, query: str, category: str = None) -> List[dict]:
|
|
"""Search memories."""
|
|
results = []
|
|
query_lower = query.lower()
|
|
base_path = self._get_category_path(category)
|
|
|
|
for md_file in base_path.rglob("*.md"):
|
|
if md_file.name.startswith("_"):
|
|
continue
|
|
content = md_file.read_text(encoding="utf-8").lower()
|
|
if query_lower in content:
|
|
fm, body = self._parse_frontmatter(md_file.read_text(encoding="utf-8"))
|
|
fm["content"] = body.strip()
|
|
fm["category"] = self._relpath(md_file.parent)
|
|
results.append(fm)
|
|
|
|
return results
|
|
|
|
def update(self, memory_id: str, category: str = None, title: str = None,
|
|
content: str = None, tags: List[str] = None, new_category: str = None) -> bool:
|
|
"""Update memory, optionally move to new category."""
|
|
old_path = self._get_memory_path(category, memory_id)
|
|
if not old_path.exists():
|
|
# Try search across all
|
|
for md_file in self.memory_dir.rglob("*.md"):
|
|
fm, _ = self._parse_frontmatter(md_file.read_text(encoding="utf-8"))
|
|
if fm.get("id") == memory_id:
|
|
old_path = md_file
|
|
category = self._relpath(md_file.parent)
|
|
break
|
|
else:
|
|
return False
|
|
|
|
fm, body = self._parse_frontmatter(old_path.read_text(encoding="utf-8"))
|
|
|
|
if title is not None:
|
|
fm['title'] = title
|
|
if content is not None:
|
|
body = content
|
|
if tags is not None:
|
|
fm['tags'] = tags
|
|
|
|
target_cat = new_category if new_category is not None else category
|
|
fm['updated'] = datetime.now().isoformat()
|
|
fm['category'] = target_cat or "/"
|
|
|
|
new_path = self._get_memory_path(target_cat, memory_id)
|
|
|
|
if new_path != old_path:
|
|
old_path.unlink()
|
|
|
|
new_path.write_text(self._create_frontmatter(fm, body), encoding="utf-8")
|
|
return True
|
|
|
|
def delete(self, memory_id: str, category: str = None) -> bool:
|
|
"""Delete memory."""
|
|
filepath = self._get_memory_path(category, memory_id)
|
|
if not filepath.exists():
|
|
# Search all
|
|
for md_file in self.memory_dir.rglob("*.md"):
|
|
fm, _ = self._parse_frontmatter(md_file.read_text(encoding="utf-8"))
|
|
if fm.get("id") == memory_id:
|
|
filepath = md_file
|
|
break
|
|
else:
|
|
return False
|
|
|
|
filepath.unlink()
|
|
|
|
# Clean empty category dirs
|
|
cat_dir = filepath.parent
|
|
if cat_dir != self.memory_dir and not any(cat_dir.iterdir()):
|
|
cat_dir.rmdir()
|
|
|
|
return True
|
|
|
|
def ls(self, category: str = None) -> dict:
|
|
"""List category structure (like ls -la)."""
|
|
cat_path = self._get_category_path(category)
|
|
result = {
|
|
"path": self._relpath(cat_path),
|
|
"categories": [],
|
|
"memories": []
|
|
}
|
|
|
|
for item in sorted(cat_path.iterdir()):
|
|
if item.name.startswith("_"):
|
|
continue
|
|
if item.is_dir():
|
|
result["categories"].append(item.name)
|
|
else:
|
|
fm, _ = self._parse_frontmatter(item.read_text(encoding="utf-8"))
|
|
result["memories"].append({
|
|
"id": fm.get("id"),
|
|
"title": fm.get("title"),
|
|
"tags": fm.get("tags", [])
|
|
})
|
|
|
|
return result
|
|
|
|
def mkdir(self, category: str) -> bool:
|
|
"""Create category."""
|
|
cat_path = self._get_category_path(category)
|
|
if cat_path.exists():
|
|
return False
|
|
cat_path.mkdir(parents=True)
|
|
return True
|
|
|
|
def rmdir(self, category: str) -> bool:
|
|
"""Remove empty category."""
|
|
cat_path = self._get_category_path(category)
|
|
if cat_path == self.memory_dir or not cat_path.exists():
|
|
return False
|
|
|
|
# Check empty
|
|
items = list(cat_path.iterdir())
|
|
if any(not i.name.startswith("_") for i in items):
|
|
return False
|
|
|
|
for item in items:
|
|
if item.is_file():
|
|
item.unlink()
|
|
elif item.is_dir():
|
|
import shutil
|
|
shutil.rmtree(item)
|
|
cat_path.rmdir()
|
|
return True
|
|
|
|
def export(self, format: str = "json", category: str = None) -> str:
|
|
"""Export memories."""
|
|
memories = self.list(category=category)
|
|
|
|
if format == "json":
|
|
return json.dumps(memories, indent=2, ensure_ascii=False)
|
|
elif format == "markdown":
|
|
lines = ["# Memory Export\n"]
|
|
current_cat = None
|
|
for m in memories:
|
|
cat = m.get("category", "/")
|
|
if cat != current_cat:
|
|
lines.append(f"\n## {cat}/\n")
|
|
current_cat = cat
|
|
lines.append(f"### {m.get('title', 'Untitled')}\n")
|
|
lines.append(f"**ID:** `{m.get('id', '')}` **Tags:** {', '.join(m.get('tags', []))}\n\n")
|
|
return "".join(lines)
|
|
elif format == "csv":
|
|
lines = ["id,title,tags,category,created,updated,content\n"]
|
|
for m in memories:
|
|
tags = "|".join(m.get('tags', []))
|
|
content = (m.get('content', '') or '').replace('"', '""').replace('\n', ' ')
|
|
lines.append(f'"{m.get("id", "")}","{m.get("title", "")}","{tags}","{m.get("category", "")}","{m.get("created", "")}","{m.get("updated", "")}","{content}"\n')
|
|
return "".join(lines)
|
|
return json.dumps(memories)
|
|
|
|
|
|
# === CLI Interface ===
|
|
|
|
def cli():
|
|
parser = argparse.ArgumentParser(description="Memory Store CLI")
|
|
parser.add_argument("--path", type=Path, default=Path(__file__).parent.parent / "memories")
|
|
|
|
subparsers = parser.add_subparsers(dest="command", help="Commands")
|
|
|
|
# Add
|
|
add_p = subparsers.add_parser("add", help="Add memory")
|
|
add_p.add_argument("-t", "--title", required=True)
|
|
add_p.add_argument("-c", "--content", required=True)
|
|
add_p.add_argument("--tags", nargs="*", default=[])
|
|
add_p.add_argument("--category", "-C", default="/")
|
|
|
|
# Get
|
|
get_p = subparsers.add_parser("get", help="Get memory")
|
|
get_p.add_argument("id")
|
|
get_p.add_argument("--category", "-C")
|
|
|
|
# List
|
|
list_p = subparsers.add_parser("list", help="List memories")
|
|
list_p.add_argument("--category", "-C")
|
|
list_p.add_argument("--no-recursive", action="store_true")
|
|
|
|
# Search
|
|
search_p = subparsers.add_parser("search", help="Search")
|
|
search_p.add_argument("query")
|
|
search_p.add_argument("--category", "-C")
|
|
|
|
# Update
|
|
update_p = subparsers.add_parser("update", help="Update")
|
|
update_p.add_argument("id")
|
|
update_p.add_argument("--category", "-C")
|
|
update_p.add_argument("-t", "--title")
|
|
update_p.add_argument("-c", "--content")
|
|
update_p.add_argument("--tags", nargs="*")
|
|
update_p.add_argument("--move", "-M", dest="new_category")
|
|
|
|
# Delete
|
|
delete_p = subparsers.add_parser("delete", help="Delete")
|
|
delete_p.add_argument("id")
|
|
delete_p.add_argument("--category", "-C")
|
|
|
|
# ls (tree)
|
|
ls_p = subparsers.add_parser("ls", help="List structure")
|
|
ls_p.add_argument("--category", "-C", default="/")
|
|
|
|
# mkdir
|
|
mkdir_p = subparsers.add_parser("mkdir", help="Create category")
|
|
mkdir_p.add_argument("category")
|
|
|
|
# rmdir
|
|
rmdir_p = subparsers.add_parser("rmdir", help="Remove empty category")
|
|
rmdir_p.add_argument("category")
|
|
|
|
# Export
|
|
export_p = subparsers.add_parser("export", help="Export")
|
|
export_p.add_argument("--format", choices=["json", "markdown", "csv"], default="json")
|
|
export_p.add_argument("-o", "--output", type=Path)
|
|
export_p.add_argument("--category", "-C")
|
|
|
|
args = parser.parse_args()
|
|
store = MemoryStore(args.path)
|
|
|
|
if args.command == "add":
|
|
mid = store.add(args.title, args.content, args.category, args.tags)
|
|
print(f"Added: {mid} -> {args.category}")
|
|
|
|
elif args.command == "get":
|
|
m = store.get(args.id, args.category)
|
|
if m:
|
|
print(f"# {m['title']}\n")
|
|
print(f"ID: {m['id']} Category: {m.get('category', '/')} Tags: {', '.join(m.get('tags', []))}")
|
|
print(f"Created: {m['created']} Updated: {m['updated']}\n")
|
|
print("---")
|
|
print(m['content'])
|
|
else:
|
|
print(f"Not found: {args.id}", file=sys.stderr)
|
|
|
|
elif args.command == "list":
|
|
for m in store.list(args.category, not args.no_recursive):
|
|
cat = m.get('category', '/')
|
|
tags = ", ".join(m.get('tags', [])[:2])
|
|
print(f"{m['id']:<10} [{cat:<15}] {m['title']:<30} [{tags}]")
|
|
|
|
elif args.command == "search":
|
|
for m in store.search(args.query, args.category):
|
|
print(f"- [{m.get('category', '/')}] {m['id']} {m['title']} ({', '.join(m.get('tags', []))})")
|
|
|
|
elif args.command == "update":
|
|
ok = store.update(args.id, args.category, args.title, args.content, args.tags, args.new_category)
|
|
print("Updated" if ok else "Not found")
|
|
|
|
elif args.command == "delete":
|
|
ok = store.delete(args.id, args.category)
|
|
print("Deleted" if ok else "Not found")
|
|
|
|
elif args.command == "ls":
|
|
result = store.ls(args.category)
|
|
print(f"📁 {result['path']}/")
|
|
if result['categories']:
|
|
for c in result['categories']:
|
|
print(f" 📂 {c}/")
|
|
for m in result['memories']:
|
|
print(f" 📄 {m['id']} - {m['title']}")
|
|
|
|
elif args.command == "mkdir":
|
|
ok = store.mkdir(args.category)
|
|
print("Created" if ok else "Already exists")
|
|
|
|
elif args.command == "rmdir":
|
|
ok = store.rmdir(args.category)
|
|
print("Removed" if ok else "Failed (not empty or not found)")
|
|
|
|
elif args.command == "export":
|
|
output = store.export(args.format, args.category)
|
|
if args.output:
|
|
args.output.write_text(output, encoding="utf-8")
|
|
print(f"Exported to {args.output}")
|
|
else:
|
|
print(output)
|
|
|
|
else:
|
|
parser.print_help()
|
|
|
|
|
|
# === MCP JSON-RPC Server ===
|
|
|
|
def mcp_server():
|
|
"""Run as MCP JSON-RPC server (stdio)."""
|
|
store = MemoryStore(Path(__file__).parent.parent / "memories")
|
|
|
|
for line in sys.stdin:
|
|
try:
|
|
request = json.loads(line.strip())
|
|
method = request.get("method", "")
|
|
params = request.get("params", {})
|
|
req_id = request.get("id")
|
|
|
|
result = None
|
|
error = None
|
|
|
|
try:
|
|
if method == "add":
|
|
result = store.add(
|
|
params.get("title", ""),
|
|
params.get("content", ""),
|
|
params.get("category"),
|
|
params.get("tags", [])
|
|
)
|
|
elif method == "get":
|
|
result = store.get(params.get("id", ""), params.get("category"))
|
|
elif method == "list":
|
|
result = store.list(params.get("category"), params.get("recursive", True))
|
|
elif method == "search":
|
|
result = store.search(params.get("query", ""), params.get("category"))
|
|
elif method == "update":
|
|
result = store.update(
|
|
params.get("id", ""),
|
|
params.get("category"),
|
|
params.get("title"),
|
|
params.get("content"),
|
|
params.get("tags"),
|
|
params.get("new_category")
|
|
)
|
|
elif method == "delete":
|
|
result = store.delete(params.get("id", ""), params.get("category"))
|
|
elif method == "ls":
|
|
result = store.ls(params.get("category"))
|
|
elif method == "mkdir":
|
|
result = store.mkdir(params.get("category", ""))
|
|
elif method == "rmdir":
|
|
result = store.rmdir(params.get("category", ""))
|
|
elif method == "export":
|
|
result = store.export(params.get("format", "json"), params.get("category"))
|
|
else:
|
|
error = {"code": -32601, "message": f"Unknown method: {method}"}
|
|
except Exception as e:
|
|
error = {"code": -32603, "message": str(e)}
|
|
|
|
response = {"jsonrpc": "2.0", "id": req_id}
|
|
if error:
|
|
response["error"] = error
|
|
else:
|
|
response["result"] = result
|
|
|
|
print(json.dumps(response), flush=True)
|
|
|
|
except json.JSONDecodeError:
|
|
continue
|
|
except Exception:
|
|
continue
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) > 1 and sys.argv[1] == "--mcp":
|
|
mcp_server()
|
|
else:
|
|
cli()
|