rex/memory.py
Rama master · réplica de lectura del repositorio original.
"""Memory management for Rex AI Assistant.
This module provides memory utilities for user profiles and conversation history,
as well as structured memory capabilities:
- Working Memory: Short-term buffer for recent interactions and task summaries
- Long-Term Memory: Structured entries with categories, expiration, and search
Re-exports from memory_utils for backward compatibility.
Ownership model (US-303 per-user isolation):
- Every working-memory and long-term-memory operation requires an explicit
validated ``user_id`` (see :func:`rex.identity.validate_user_id`).
- Stores are partitioned on disk per user::
data/memory/<user_id>/working_memory.json
data/memory/<user_id>/long_term_memory.json
- Missing, blank, malformed, or traversal-style identity fails closed
(``TypeError``/``ValueError``); there is no silent fallback to ``default``,
the active user, or any other profile.
- Process-level instances are cached in dictionaries keyed by validated
``user_id``; one user never receives an object backed by another user's
path or in-memory entries.
- User IDs that differ only by case from an existing profile are rejected
on every platform: Windows/macOS filesystems are case-insensitive, so
"James" and "james" would otherwise silently alias one on-disk store.
Legacy unscoped-file migration:
- The pre-isolation shared files ``data/memory/working_memory.json`` and
``data/memory/long_term_memory.json`` belong only to the distinct
``default`` profile. They are never shared and never reassigned to a named
user, and named users never read them.
- Migration runs only when a caller explicitly requests
``user_id="default"``. It is idempotent and crash-safe: the legacy file is
copied into ``data/memory/default/`` via a temp file + atomic replace, and
only after the default-profile copy exists is the original moved aside as
``<name>.json.pre-user-isolation.bak`` (never deleted). A failed or
partial migration leaves the original untouched and is retried on the next
explicit ``default`` access. To recover manually, restore the ``.bak``
file to its original name and remove ``data/memory/default/<name>.json``.
Usage:
from rex.memory import (
get_working_memory,
get_long_term_memory,
add_user_preference,
remember_context,
)
# Working memory (owner-scoped)
wm = get_working_memory(user_id="james")
wm.add_entry("User asked about weather")
# Long-term memory (owner-scoped)
ltm = get_long_term_memory(user_id="james")
entry = ltm.add_entry("preferences", {"theme": "dark"})
"""
from __future__ import annotations
import json
import logging
import os
import shutil
import uuid
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import Any
from pydantic import BaseModel, Field
from rex.identity import validate_user_id
from rex.runtime_paths import data_dir, user_data_dir, users_data_dir
# Re-export existing memory utilities for backward compatibility
from .memory_utils import ( # noqa: F401
append_history_entry,
export_transcript,
extract_voice_reference,
load_all_profiles,
load_memory_profile,
load_recent_history,
load_users_map,
resolve_user_key,
trim_history,
)
logger = logging.getLogger(__name__)
# Compatibility seam for callers/tests that replace the structured-memory
# root. Production defaults use the canonical private-user layout.
_DATA_DIR = data_dir() / "memory"
_INITIAL_DATA_DIR = _DATA_DIR
def _legacy_memory_dir() -> Path:
"""Return the pre-canonical structured-memory root."""
return _DATA_DIR
# The distinct profile that owns pre-isolation legacy data. Never used as an
# implicit fallback; callers must select it explicitly.
DEFAULT_PROFILE = "default"
# Suffix appended to a legacy shared store file once its content has been
# migrated into the explicit ``default`` profile.
LEGACY_BACKUP_SUFFIX = ".pre-user-isolation.bak"
_WORKING_MEMORY_FILENAME = "working_memory.json"
_LONG_TERM_MEMORY_FILENAME = "long_term_memory.json"
def _utc_now() -> datetime:
"""Return the current UTC datetime."""
return datetime.now(UTC)
def _reject_case_conflicts(user_id: str) -> None:
"""Fail closed when *user_id* differs only by case from a known profile.
Windows (NTFS) and macOS (APFS/HFS+) filesystems are case-insensitive,
so two IDs differing only by case would silently alias one on-disk
store — user "James" would read and rewrite user "james"'s files. The
later spelling is rejected on every platform so behavior is consistent
and fail-closed.
Raises:
ValueError: If a case-variant of *user_id* already exists in the
process registries or as a memory directory on disk.
"""
folded = user_id.casefold()
known_ids: set[str] = set(_working_memories) | set(_long_term_memories)
private_root = users_data_dir()
if private_root.is_dir():
known_ids.update(entry.name for entry in private_root.iterdir() if entry.is_dir())
legacy_root = _legacy_memory_dir()
if legacy_root.is_dir():
known_ids.update(entry.name for entry in legacy_root.iterdir() if entry.is_dir())
for existing in known_ids:
if existing != user_id and existing.casefold() == folded:
raise ValueError(
f"Invalid user_id: {user_id!r} (conflicts with an existing memory profile)"
)
def _user_memory_dir(user_id: str) -> Path:
"""Return private memory storage and copy a legacy user store on first access."""
user_id = validate_user_id(user_id)
_reject_case_conflicts(user_id)
target = (
_DATA_DIR / user_id if _DATA_DIR != _INITIAL_DATA_DIR else user_data_dir(user_id) / "memory"
)
legacy = _legacy_memory_dir() / user_id
if legacy.is_dir() and not target.exists():
target.parent.mkdir(parents=True, exist_ok=True)
shutil.copytree(legacy, target)
logger.info("Copied legacy memory profile '%s' into canonical private storage", user_id)
return target
def _migrate_legacy_default_store(filename: str) -> None:
"""Migrate a pre-isolation shared store file to the ``default`` profile.
Called only when a caller explicitly requests ``user_id="default"``.
Idempotent and crash-safe:
1. If ``data/memory/<filename>`` does not exist, do nothing.
2. If the default-profile copy does not exist yet, copy the legacy file
via a temp file and atomic replace. A failure here propagates (the
caller fails closed), leaves the original untouched, and is retried
on the next explicit ``default`` access.
3. Only once the default-profile copy exists is the original moved aside
as ``<filename>.pre-user-isolation.bak`` — preserved, never deleted.
If the move fails, the original stays in place and the move is
retried later; migration itself is already complete.
"""
legacy_root = _legacy_memory_dir()
legacy = legacy_root / filename
if not legacy.is_file():
return
target = _user_memory_dir(DEFAULT_PROFILE) / filename
if not target.exists():
target.parent.mkdir(parents=True, exist_ok=True)
# PID-unique temp file: concurrent first-time migrations (e.g. two
# GUI bridge processes) each copy privately, and os.replace is
# atomic, so interleaved writes can never publish a corrupt store.
tmp_copy = target.with_name(f"{target.name}.migrating.{os.getpid()}")
try:
shutil.copyfile(legacy, tmp_copy)
os.replace(tmp_copy, target)
finally:
tmp_copy.unlink(missing_ok=True)
logger.info("Migrated legacy shared %s into the 'default' profile store", filename)
backup = legacy_root / (filename + LEGACY_BACKUP_SUFFIX)
if not backup.exists():
try:
legacy.rename(backup)
logger.info("Preserved legacy %s as %s", filename, backup.name)
except OSError as exc:
logger.warning(
"Could not move legacy %s aside: %s (original left in place)",
filename,
exc,
)
# =============================================================================
# Working Memory
# =============================================================================
class WorkingMemory:
"""Short-term memory buffer for recent interactions and task summaries.
Holds an ordered list of recent entries, automatically persisting to disk
and loading on startup. Useful for maintaining immediate conversational
context. Each instance is backed by exactly one owner's store.
Attributes:
max_entries: Maximum number of entries to retain (default 100).
storage_path: Path to the persistence file.
user_id: Validated owner of this store, or None for an explicit
custom ``storage_path`` (unit-test instances).
Example:
wm = WorkingMemory(user_id="james")
wm.add_entry("User asked about the weather in Dallas")
wm.add_entry("Checked weather API - sunny, 72°F")
recent = wm.get_recent(5)
for entry in recent:
print(entry)
"""
def __init__(
self,
storage_path: Path | str | None = None,
max_entries: int = 100,
*,
user_id: str | None = None,
) -> None:
"""Initialize working memory for one owner.
Args:
storage_path: Explicit persistence file (unit-test support). When
omitted, a validated ``user_id`` is required and the store
lives at ``data/memory/<user_id>/working_memory.json``.
max_entries: Maximum entries to retain before oldest are removed.
user_id: Owner of this store. Validated; fails closed when the
default per-user path is used and no owner is given.
Raises:
ValueError: If neither ``storage_path`` nor a valid ``user_id``
is provided, or if ``user_id`` is invalid.
"""
if user_id is not None:
user_id = validate_user_id(user_id)
if storage_path is None:
if user_id is None:
raise ValueError("WorkingMemory requires a user_id (or an explicit storage_path)")
if user_id == DEFAULT_PROFILE:
_migrate_legacy_default_store(_WORKING_MEMORY_FILENAME)
storage_path = _user_memory_dir(user_id) / _WORKING_MEMORY_FILENAME
self.user_id = user_id
self.storage_path = Path(storage_path)
self.max_entries = max_entries
self._entries: list[dict[str, Any]] = []
self._load()
def _owner_label(self) -> str:
"""Owner identifier for logging (never memory content)."""
return self.user_id if self.user_id is not None else "<custom-path>"
def _load(self) -> None:
"""Load entries from disk."""
if self.storage_path.exists():
try:
with open(self.storage_path, encoding="utf-8") as f:
data = json.load(f)
self._entries = data.get("entries", [])
logger.debug(
"Loaded %d working memory entries for owner '%s'",
len(self._entries),
self._owner_label(),
)
except (json.JSONDecodeError, OSError) as e:
logger.warning(
"Failed to load working memory for owner '%s': %s",
self._owner_label(),
e,
)
self._entries = []
def _save(self) -> None:
"""Save entries to disk."""
self.storage_path.parent.mkdir(parents=True, exist_ok=True)
try:
with open(self.storage_path, "w", encoding="utf-8") as f:
json.dump({"entries": self._entries}, f, indent=2, default=str)
except OSError as e:
logger.error(
"Failed to save working memory for owner '%s': %s",
self._owner_label(),
e,
)
def add_entry(self, content: str) -> None:
"""Add a new entry to working memory.
Args:
content: The content to add (interaction summary, note, etc.)
"""
entry = {
"content": content,
"timestamp": _utc_now().isoformat(),
}
self._entries.append(entry)
# Trim to max entries
if len(self._entries) > self.max_entries:
self._entries = self._entries[-self.max_entries :]
self._save()
def get_recent(self, n: int = 10) -> list[str]:
"""Get the most recent entries.
Args:
n: Number of entries to retrieve.
Returns:
List of content strings, most recent last.
"""
recent = self._entries[-n:] if n < len(self._entries) else self._entries
return [entry["content"] for entry in recent]
def get_recent_with_timestamps(self, n: int = 10) -> list[dict[str, Any]]:
"""Get recent entries with their timestamps.
Args:
n: Number of entries to retrieve.
Returns:
List of entry dicts with 'content' and 'timestamp' keys.
"""
return self._entries[-n:] if n < len(self._entries) else self._entries.copy()
def clear(self) -> None:
"""Clear all working memory entries for this owner."""
self._entries = []
self._save()
logger.debug("Cleared working memory for owner '%s'", self._owner_label())
def __len__(self) -> int:
"""Return the number of entries."""
return len(self._entries)
def stats(self) -> dict[str, Any]:
"""Return summary statistics for working memory."""
return {
"entries": len(self._entries),
"max_entries": self.max_entries,
}
# =============================================================================
# Long-Term Memory
# =============================================================================
class MemoryEntry(BaseModel):
"""A structured long-term memory entry.
Attributes:
entry_id: Unique identifier for this entry.
category: Category for organizing entries (e.g., 'preferences', 'facts').
content: The stored data as a dictionary.
created_at: When the entry was created.
expires_at: When the entry should expire (None for no expiration).
sensitive: Whether this entry contains sensitive data.
"""
entry_id: str = Field(
default_factory=lambda: f"mem_{uuid.uuid4().hex[:12]}",
description="Unique identifier for this entry",
)
category: str = Field(
...,
description="Category for organizing entries",
)
content: dict[str, Any] = Field(
default_factory=dict,
description="The stored data",
)
created_at: datetime = Field(
default_factory=_utc_now,
description="When the entry was created",
)
expires_at: datetime | None = Field(
default=None,
description="When the entry should expire (None for no expiration)",
)
sensitive: bool = Field(
default=False,
description="Whether this entry contains sensitive data",
)
def is_expired(self) -> bool:
"""Check if this entry has expired."""
if self.expires_at is None:
return False
return _utc_now() > self.expires_at
def to_safe_dict(self) -> dict[str, Any]:
"""Return a dictionary with sensitive content redacted."""
data = self.model_dump()
if self.sensitive:
data["content"] = {"[SENSITIVE]": "Content hidden"}
return data
model_config = {
"json_schema_extra": {
"examples": [
{
"entry_id": "mem_abc123def456",
"category": "preferences",
"content": {"theme": "dark", "language": "en"},
"created_at": "2024-01-15T10:30:00Z",
"expires_at": None,
"sensitive": False,
}
]
}
}
class LongTermMemory:
"""Long-term structured memory with retention policies.
Stores entries organized by category with support for:
- Expiration dates and automatic cleanup
- Sensitive data flagging
- Keyword-based search across content
Entries are persisted to disk as JSON. Each instance is backed by exactly
one owner's store; entry IDs may repeat across owners without collision.
Example:
ltm = LongTermMemory(user_id="james")
# Add a preference that expires in 30 days
entry = ltm.add_entry(
category="preferences",
content={"theme": "dark"},
expires_in=timedelta(days=30),
)
# Search for entries
results = ltm.search(keyword="theme")
# Clean up expired entries
ltm.run_retention_policy()
"""
def __init__(
self,
storage_path: Path | str | None = None,
*,
user_id: str | None = None,
) -> None:
"""Initialize long-term memory for one owner.
Args:
storage_path: Explicit persistence file (unit-test support). When
omitted, a validated ``user_id`` is required and the store
lives at ``data/memory/<user_id>/long_term_memory.json``.
user_id: Owner of this store. Validated; fails closed when the
default per-user path is used and no owner is given.
Raises:
ValueError: If neither ``storage_path`` nor a valid ``user_id``
is provided, or if ``user_id`` is invalid.
"""
if user_id is not None:
user_id = validate_user_id(user_id)
if storage_path is None:
if user_id is None:
raise ValueError("LongTermMemory requires a user_id (or an explicit storage_path)")
if user_id == DEFAULT_PROFILE:
_migrate_legacy_default_store(_LONG_TERM_MEMORY_FILENAME)
storage_path = _user_memory_dir(user_id) / _LONG_TERM_MEMORY_FILENAME
self.user_id = user_id
self.storage_path = Path(storage_path)
self._entries: dict[str, MemoryEntry] = {}
self._load()
self.run_retention_policy()
def _owner_label(self) -> str:
"""Owner identifier for logging (never memory content)."""
return self.user_id if self.user_id is not None else "<custom-path>"
def _load(self) -> None:
"""Load entries from disk."""
if self.storage_path.exists():
try:
with open(self.storage_path, encoding="utf-8") as f:
data = json.load(f)
for entry_data in data.get("entries", []):
entry = MemoryEntry.model_validate(entry_data)
self._entries[entry.entry_id] = entry
logger.debug(
"Loaded %d long-term memory entries for owner '%s'",
len(self._entries),
self._owner_label(),
)
except (json.JSONDecodeError, OSError) as e:
logger.warning(
"Failed to load long-term memory for owner '%s': %s",
self._owner_label(),
e,
)
self._entries = {}
def _save(self) -> None:
"""Save entries to disk."""
self.storage_path.parent.mkdir(parents=True, exist_ok=True)
try:
entries_data = [entry.model_dump(mode="json") for entry in self._entries.values()]
with open(self.storage_path, "w", encoding="utf-8") as f:
json.dump({"entries": entries_data}, f, indent=2, default=str)
except OSError as e:
logger.error(
"Failed to save long-term memory for owner '%s': %s",
self._owner_label(),
e,
)
def add_entry(
self,
category: str,
content: dict[str, Any],
expires_in: timedelta | None = None,
sensitive: bool = False,
entry_id: str | None = None,
) -> MemoryEntry:
"""Add a new entry to long-term memory.
Args:
category: Category for the entry (e.g., 'preferences', 'facts').
content: The data to store as a dictionary.
expires_in: Time until expiration (None for no expiration).
sensitive: Whether this contains sensitive data.
entry_id: Optional custom entry ID.
Returns:
The created MemoryEntry.
"""
expires_at = None
if expires_in is not None:
expires_at = _utc_now() + expires_in
entry = MemoryEntry(
entry_id=entry_id or f"mem_{uuid.uuid4().hex[:12]}",
category=category,
content=content,
expires_at=expires_at,
sensitive=sensitive,
)
self._entries[entry.entry_id] = entry
self._save()
logger.debug(
"Added memory entry %s in category '%s' for owner '%s'",
entry.entry_id,
category,
self._owner_label(),
)
return entry
def get_entry(self, entry_id: str) -> MemoryEntry | None:
"""Get a specific entry by ID.
Args:
entry_id: The entry ID to look up.
Returns:
The MemoryEntry if found and not expired, else None.
"""
entry = self._entries.get(entry_id)
if entry is None or entry.is_expired():
return None
return entry
def search(
self,
category: str | None = None,
keyword: str | None = None,
include_sensitive: bool = True,
include_expired: bool = False,
) -> list[MemoryEntry]:
"""Search for entries by category and/or keyword.
Args:
category: Filter by category (exact match).
keyword: Search for keyword in content keys and values (substring).
include_sensitive: Whether to include sensitive entries.
include_expired: Whether to include expired entries.
Returns:
List of matching MemoryEntry objects.
"""
results = []
keyword_lower = keyword.lower() if keyword else None
for entry in self._entries.values():
# Skip expired entries unless requested
if not include_expired and entry.is_expired():
continue
# Skip sensitive entries if not requested
if not include_sensitive and entry.sensitive:
continue
# Category filter
if category is not None and entry.category != category:
continue
# Keyword filter - search in content keys and values
if keyword_lower is not None:
if not self._matches_keyword(entry.content, keyword_lower):
# Also check category
if keyword_lower not in entry.category.lower():
continue
results.append(entry)
# Sort by created_at descending (newest first)
results.sort(key=lambda e: e.created_at, reverse=True)
return results
def _matches_keyword(self, content: dict[str, Any], keyword: str) -> bool:
"""Check if content matches keyword (recursive substring search)."""
def search_value(value: Any) -> bool:
if isinstance(value, str):
return keyword in value.lower()
elif isinstance(value, dict):
for k, v in value.items():
if keyword in str(k).lower():
return True
if search_value(v):
return True
elif isinstance(value, (list, tuple)):
for item in value:
if search_value(item):
return True
else:
return keyword in str(value).lower()
return False
return search_value(content)
def forget(self, entry_id: str) -> bool:
"""Delete a specific entry.
Args:
entry_id: The entry ID to delete.
Returns:
True if the entry was deleted, False if not found.
"""
if entry_id in self._entries:
del self._entries[entry_id]
self._save()
logger.debug("Deleted memory entry %s for owner '%s'", entry_id, self._owner_label())
return True
return False
def run_retention_policy(self) -> int:
"""Delete all expired entries.
This is automatically called on startup and can be called manually.
Returns:
Number of entries deleted.
"""
expired_ids = [entry_id for entry_id, entry in self._entries.items() if entry.is_expired()]
for entry_id in expired_ids:
del self._entries[entry_id]
if expired_ids:
self._save()
logger.info(
"Retention policy deleted %d expired entries for owner '%s'",
len(expired_ids),
self._owner_label(),
)
return len(expired_ids)
def compact(self) -> int:
"""Remove expired entries and rewrite storage to reclaim space.
Unlike run_retention_policy which only saves when entries were deleted,
compact always rewrites the storage file to ensure it contains only
current entries with no stale data.
Returns:
Number of expired entries removed.
"""
removed = self.run_retention_policy()
# Always rewrite storage to compact the file
self._save()
logger.info(
"Memory store compacted for owner '%s': %d expired entries removed, "
"%d active entries retained",
self._owner_label(),
removed,
len(self),
)
return removed
def list_categories(self) -> list[str]:
"""List all unique categories.
Returns:
Sorted list of category names.
"""
categories = set()
for entry in self._entries.values():
if not entry.is_expired():
categories.add(entry.category)
return sorted(categories)
def count_by_category(self) -> dict[str, int]:
"""Count non-expired entries per category.
Returns:
Dictionary of category name to count.
"""
counts: dict[str, int] = {}
for entry in self._entries.values():
if not entry.is_expired():
counts[entry.category] = counts.get(entry.category, 0) + 1
return counts
def __len__(self) -> int:
"""Return the number of non-expired entries."""
return sum(1 for e in self._entries.values() if not e.is_expired())
def stats(self) -> dict[str, Any]:
"""Return summary statistics for long-term memory."""
return {
"entries": len(self),
"categories": self.count_by_category(),
}
# =============================================================================
# Per-User Instance Registries
# =============================================================================
_working_memories: dict[str, WorkingMemory] = {}
_long_term_memories: dict[str, LongTermMemory] = {}
def get_working_memory(*, user_id: str) -> WorkingMemory:
"""Get the working memory instance owned by *user_id*.
Args:
user_id: Validated owner of the store. Required; missing or invalid
identity fails closed (``TypeError``/``ValueError``).
Returns:
The per-user WorkingMemory instance (created lazily).
"""
user_id = validate_user_id(user_id)
instance = _working_memories.get(user_id)
if instance is None:
instance = WorkingMemory(user_id=user_id)
_working_memories[user_id] = instance
return instance
def set_working_memory(wm: WorkingMemory | None, *, user_id: str) -> None:
"""Set (or clear) the working memory instance for *user_id* (for testing)."""
user_id = validate_user_id(user_id)
if wm is None:
_working_memories.pop(user_id, None)
else:
_reject_case_conflicts(user_id)
_working_memories[user_id] = wm
def get_long_term_memory(*, user_id: str) -> LongTermMemory:
"""Get the long-term memory instance owned by *user_id*.
Args:
user_id: Validated owner of the store. Required; missing or invalid
identity fails closed (``TypeError``/``ValueError``).
Returns:
The per-user LongTermMemory instance (created lazily).
"""
user_id = validate_user_id(user_id)
instance = _long_term_memories.get(user_id)
if instance is None:
instance = LongTermMemory(user_id=user_id)
_long_term_memories[user_id] = instance
return instance
def set_long_term_memory(ltm: LongTermMemory | None, *, user_id: str) -> None:
"""Set (or clear) the long-term memory instance for *user_id* (for testing)."""
user_id = validate_user_id(user_id)
if ltm is None:
_long_term_memories.pop(user_id, None)
else:
_reject_case_conflicts(user_id)
_long_term_memories[user_id] = ltm
# =============================================================================
# Convenience Functions
# =============================================================================
def add_user_preference(
key: str,
value: Any,
expires_in: timedelta | None = None,
sensitive: bool = False,
*,
user_id: str,
) -> MemoryEntry:
"""Add a user preference to *user_id*'s long-term memory.
Args:
key: Preference key (e.g., 'theme', 'language').
value: Preference value.
expires_in: Optional expiration time.
sensitive: Whether this is sensitive data.
user_id: Owner of the preference. Required; fails closed on
missing or invalid identity.
Returns:
The created MemoryEntry.
"""
ltm = get_long_term_memory(user_id=user_id)
return ltm.add_entry(
category="user_preferences",
content={key: value},
expires_in=expires_in,
sensitive=sensitive,
)
def get_user_preferences(key: str | None = None, *, user_id: str) -> list[MemoryEntry]:
"""Get *user_id*'s preferences from long-term memory.
Args:
key: Optional key to search for.
user_id: Owner whose preferences to read. Required; fails closed on
missing or invalid identity.
Returns:
List of matching preference entries.
"""
ltm = get_long_term_memory(user_id=user_id)
return ltm.search(category="user_preferences", keyword=key)
def add_fact(
topic: str,
content: dict[str, Any],
expires_in: timedelta | None = None,
*,
user_id: str,
) -> MemoryEntry:
"""Add a fact to *user_id*'s long-term memory.
Args:
topic: The topic or subject of the fact.
content: Fact data.
expires_in: Optional expiration time.
user_id: Owner of the fact. Required; fails closed on missing or
invalid identity.
Returns:
The created MemoryEntry.
"""
ltm = get_long_term_memory(user_id=user_id)
return ltm.add_entry(
category="facts",
content={"topic": topic, **content},
expires_in=expires_in,
)
def remember_context(summary: str, *, user_id: str) -> None:
"""Add a context summary to *user_id*'s working memory.
Args:
summary: A summary of the current context or interaction.
user_id: Owner of the context. Required; fails closed on missing or
invalid identity.
"""
wm = get_working_memory(user_id=user_id)
wm.add_entry(summary)
def get_recent_context(n: int = 5, *, user_id: str) -> list[str]:
"""Get recent context from *user_id*'s working memory.
Args:
n: Number of entries to retrieve.
user_id: Owner whose context to read. Required; fails closed on
missing or invalid identity.
Returns:
List of recent context summaries.
"""
wm = get_working_memory(user_id=user_id)
return wm.get_recent(n)
# =============================================================================
# Per-User Maintenance (cleanup scheduling and metrics)
# =============================================================================
def list_memory_user_ids() -> list[str]:
"""Discover user IDs that have a per-user memory directory on disk.
Only directory names that pass :func:`rex.identity.validate_user_id` are
returned; anything else (including the legacy shared files and their
backups at the data-dir root) is ignored and never treated as a user.
"""
root = _DATA_DIR if _DATA_DIR != _INITIAL_DATA_DIR else users_data_dir()
if not root.is_dir():
return []
users: list[str] = []
for entry in sorted(root.iterdir()):
if not entry.is_dir():
continue
try:
validate_user_id(entry.name)
except ValueError:
logger.warning("Ignoring invalid memory directory name: %r", entry.name)
continue
users.append(entry.name)
return users
def run_memory_cleanup() -> dict[str, int]:
"""Compact every known user's long-term store, each independently.
Covers users discovered on disk plus any live registry instances (e.g.
test-injected stores). A failure in one user's store is logged and never
affects another user's store.
Returns:
Mapping of user_id to the number of expired entries removed. Users
whose cleanup failed are omitted.
"""
user_ids = sorted(set(list_memory_user_ids()) | set(_long_term_memories))
results: dict[str, int] = {}
for user_id in user_ids:
try:
removed = get_long_term_memory(user_id=user_id).compact()
except Exception:
logger.exception(
"Memory cleanup failed for owner '%s'; other stores unaffected", user_id
)
continue
results[user_id] = removed
return results
def schedule_memory_cleanup(
scheduler: Any,
interval_seconds: int = 3600,
job_id: str = "memory_cleanup",
) -> None:
"""Register a scheduled memory cleanup job with the scheduler.
Registers a callback that runs :func:`run_memory_cleanup` at the given
interval: every user's long-term store is compacted independently, one
user's failure never touches another user's file, and invalid directory
names are ignored.
Args:
scheduler: A Scheduler instance from rex.scheduler.
interval_seconds: How often to run cleanup (default: 3600 = 1 hour).
job_id: Unique ID for the scheduler job.
"""
def _cleanup_callback(job: Any) -> None: # noqa: ARG001
results = run_memory_cleanup()
removed = sum(results.values())
logger.info(
"Scheduled memory cleanup complete: %d expired entries removed "
"across %d user store(s)",
removed,
len(results),
)
# Register callback and add the job
scheduler.register_callback("memory_cleanup", _cleanup_callback)
scheduler.add_job(
job_id=job_id,
name="Memory Cleanup",
schedule=f"interval:{interval_seconds}",
callback_name="memory_cleanup",
)
logger.info(f"Memory cleanup scheduled every {interval_seconds}s (job_id={job_id})")
def memory_store_metrics() -> dict[str, Any]:
"""Aggregate entry counts across all per-user stores (no content).
Used by the service supervisor for health metrics. Counts only — never
memory content, categories, or entry IDs.
"""
user_ids = sorted(
set(list_memory_user_ids()) | set(_working_memories) | set(_long_term_memories)
)
working_total = 0
long_term_total = 0
for user_id in user_ids:
try:
working_total += len(get_working_memory(user_id=user_id))
long_term_total += len(get_long_term_memory(user_id=user_id))
except Exception:
logger.exception("Failed to collect memory metrics for owner '%s'", user_id)
return {
"user_count": len(user_ids),
"working_memory_entries": working_total,
"long_term_memory_entries": long_term_total,
}
# Update __all__ to include new exports
__all__ = [
# Legacy memory utilities
"load_users_map",
"resolve_user_key",
"load_memory_profile",
"load_all_profiles",
"extract_voice_reference",
"trim_history",
"append_history_entry",
"load_recent_history",
"export_transcript",
# Ownership constants
"DEFAULT_PROFILE",
"LEGACY_BACKUP_SUFFIX",
# Working memory
"WorkingMemory",
"get_working_memory",
"set_working_memory",
# Long-term memory
"MemoryEntry",
"LongTermMemory",
"get_long_term_memory",
"set_long_term_memory",
# Convenience functions
"add_user_preference",
"get_user_preferences",
"add_fact",
"remember_context",
"get_recent_context",
# Per-user maintenance
"list_memory_user_ids",
"run_memory_cleanup",
"schedule_memory_cleanup",
"memory_store_metrics",
]