Teknium
1452c81941
Introduces MemoryProvider ABC, MemoryManager orchestrator, and BuiltinMemoryProvider for the existing MEMORY.md/USER.md system. Key design decisions: - Built-in memory is ALWAYS active, never disabled by external providers - Multiple providers can be active simultaneously - Prefetch results from all providers are merged per-turn - Sync fans out to all providers after each turn - Each provider can expose its own tools Three registration paths: 1. Built-in (BuiltinMemoryProvider) — always first, not removable 2. First-party (Honcho stays as-is for now, migration in follow-up) 3. Plugin — ctx.register_memory_provider() in plugin system Files: - agent/memory_provider.py — ABC with core + optional lifecycle hooks - agent/memory_manager.py — orchestrator, single integration point - agent/builtin_memory_provider.py — wraps existing MemoryStore - hermes_cli/plugins.py — register_memory_provider() + accessor - run_agent.py — MemoryManager wired alongside existing Honcho code - tests/agent/test_memory_provider.py — 37 tests This establishes the interface for all pending memory backend PRs (#1811 Hindsight, #2732 RetainDB, #2933 Mem0, #3499 Byterover, #3369 OpenViking, #2351 Holographic, #727 Cognitive) to implement as plugins rather than one-off integrations.
282 lines
10 KiB
Python
282 lines
10 KiB
Python
"""MemoryManager — orchestrates multiple memory providers.
|
|
|
|
Single integration point in run_agent.py. Replaces scattered per-backend
|
|
code with one manager that delegates to all registered providers.
|
|
|
|
The BuiltinMemoryProvider is always registered first and cannot be removed.
|
|
External providers are additive — they never disable the built-in store.
|
|
|
|
Usage in run_agent.py:
|
|
self._memory_manager = MemoryManager()
|
|
self._memory_manager.add_provider(BuiltinMemoryProvider(...))
|
|
if honcho_configured:
|
|
self._memory_manager.add_provider(HonchoProvider(...))
|
|
# Plugin providers are added via register_memory_provider()
|
|
|
|
# System prompt
|
|
prompt_parts.append(self._memory_manager.build_system_prompt())
|
|
|
|
# Pre-turn
|
|
context = self._memory_manager.prefetch_all(user_message)
|
|
|
|
# Post-turn
|
|
self._memory_manager.sync_all(user_msg, assistant_response)
|
|
self._memory_manager.queue_prefetch_all(user_msg)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from agent.memory_provider import MemoryProvider
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class MemoryManager:
|
|
"""Orchestrates multiple memory providers.
|
|
|
|
Providers are called in registration order. The builtin provider
|
|
is always first. Failures in one provider never block others.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self._providers: List[MemoryProvider] = []
|
|
self._tool_to_provider: Dict[str, MemoryProvider] = {}
|
|
|
|
# -- Registration --------------------------------------------------------
|
|
|
|
def add_provider(self, provider: MemoryProvider) -> None:
|
|
"""Register a memory provider.
|
|
|
|
Providers are called in registration order for all operations.
|
|
Tool name conflicts are resolved first-registered-wins.
|
|
"""
|
|
self._providers.append(provider)
|
|
|
|
# Index tool names → provider for routing
|
|
for schema in provider.get_tool_schemas():
|
|
tool_name = schema.get("name", "")
|
|
if tool_name and tool_name not in self._tool_to_provider:
|
|
self._tool_to_provider[tool_name] = provider
|
|
elif tool_name in self._tool_to_provider:
|
|
logger.warning(
|
|
"Memory tool name conflict: '%s' already registered by %s, "
|
|
"ignoring from %s",
|
|
tool_name,
|
|
self._tool_to_provider[tool_name].name,
|
|
provider.name,
|
|
)
|
|
|
|
logger.info(
|
|
"Memory provider '%s' registered (%d tools)",
|
|
provider.name,
|
|
len(provider.get_tool_schemas()),
|
|
)
|
|
|
|
@property
|
|
def providers(self) -> List[MemoryProvider]:
|
|
"""All registered providers in order."""
|
|
return list(self._providers)
|
|
|
|
@property
|
|
def provider_names(self) -> List[str]:
|
|
"""Names of all registered providers."""
|
|
return [p.name for p in self._providers]
|
|
|
|
def get_provider(self, name: str) -> Optional[MemoryProvider]:
|
|
"""Get a provider by name, or None if not registered."""
|
|
for p in self._providers:
|
|
if p.name == name:
|
|
return p
|
|
return None
|
|
|
|
# -- System prompt -------------------------------------------------------
|
|
|
|
def build_system_prompt(self) -> str:
|
|
"""Collect system prompt blocks from all providers.
|
|
|
|
Returns combined text, or empty string if no providers contribute.
|
|
Each non-empty block is labeled with the provider name.
|
|
"""
|
|
blocks = []
|
|
for provider in self._providers:
|
|
try:
|
|
block = provider.system_prompt_block()
|
|
if block and block.strip():
|
|
blocks.append(block)
|
|
except Exception as e:
|
|
logger.warning(
|
|
"Memory provider '%s' system_prompt_block() failed: %s",
|
|
provider.name, e,
|
|
)
|
|
return "\n\n".join(blocks)
|
|
|
|
# -- Prefetch / recall ---------------------------------------------------
|
|
|
|
def prefetch_all(self, query: str) -> str:
|
|
"""Collect prefetch context from all providers.
|
|
|
|
Returns merged context text labeled by provider. Empty providers
|
|
are skipped. Failures in one provider don't block others.
|
|
"""
|
|
parts = []
|
|
for provider in self._providers:
|
|
try:
|
|
result = provider.prefetch(query)
|
|
if result and result.strip():
|
|
parts.append(result)
|
|
except Exception as e:
|
|
logger.debug(
|
|
"Memory provider '%s' prefetch failed (non-fatal): %s",
|
|
provider.name, e,
|
|
)
|
|
return "\n\n".join(parts)
|
|
|
|
def queue_prefetch_all(self, query: str) -> None:
|
|
"""Queue background prefetch on all providers for the next turn."""
|
|
for provider in self._providers:
|
|
try:
|
|
provider.queue_prefetch(query)
|
|
except Exception as e:
|
|
logger.debug(
|
|
"Memory provider '%s' queue_prefetch failed (non-fatal): %s",
|
|
provider.name, e,
|
|
)
|
|
|
|
# -- Sync ----------------------------------------------------------------
|
|
|
|
def sync_all(self, user_content: str, assistant_content: str) -> None:
|
|
"""Sync a completed turn to all providers."""
|
|
for provider in self._providers:
|
|
try:
|
|
provider.sync_turn(user_content, assistant_content)
|
|
except Exception as e:
|
|
logger.warning(
|
|
"Memory provider '%s' sync_turn failed: %s",
|
|
provider.name, e,
|
|
)
|
|
|
|
# -- Tools ---------------------------------------------------------------
|
|
|
|
def get_all_tool_schemas(self) -> List[Dict[str, Any]]:
|
|
"""Collect tool schemas from all providers."""
|
|
schemas = []
|
|
seen = set()
|
|
for provider in self._providers:
|
|
try:
|
|
for schema in provider.get_tool_schemas():
|
|
name = schema.get("name", "")
|
|
if name and name not in seen:
|
|
schemas.append(schema)
|
|
seen.add(name)
|
|
except Exception as e:
|
|
logger.warning(
|
|
"Memory provider '%s' get_tool_schemas() failed: %s",
|
|
provider.name, e,
|
|
)
|
|
return schemas
|
|
|
|
def get_all_tool_names(self) -> set:
|
|
"""Return set of all tool names across all providers."""
|
|
return set(self._tool_to_provider.keys())
|
|
|
|
def has_tool(self, tool_name: str) -> bool:
|
|
"""Check if any provider handles this tool."""
|
|
return tool_name in self._tool_to_provider
|
|
|
|
def handle_tool_call(
|
|
self, tool_name: str, args: Dict[str, Any], **kwargs
|
|
) -> str:
|
|
"""Route a tool call to the correct provider.
|
|
|
|
Returns JSON string result. Raises ValueError if no provider
|
|
handles the tool.
|
|
"""
|
|
provider = self._tool_to_provider.get(tool_name)
|
|
if provider is None:
|
|
return json.dumps({"error": f"No memory provider handles tool '{tool_name}'"})
|
|
try:
|
|
return provider.handle_tool_call(tool_name, args, **kwargs)
|
|
except Exception as e:
|
|
logger.error(
|
|
"Memory provider '%s' handle_tool_call(%s) failed: %s",
|
|
provider.name, tool_name, e,
|
|
)
|
|
return json.dumps({"error": f"Memory tool '{tool_name}' failed: {e}"})
|
|
|
|
# -- Lifecycle hooks -----------------------------------------------------
|
|
|
|
def on_turn_start(self, turn_number: int, message: str) -> None:
|
|
"""Notify all providers of a new turn."""
|
|
for provider in self._providers:
|
|
try:
|
|
provider.on_turn_start(turn_number, message)
|
|
except Exception as e:
|
|
logger.debug(
|
|
"Memory provider '%s' on_turn_start failed: %s",
|
|
provider.name, e,
|
|
)
|
|
|
|
def on_session_end(self, messages: List[Dict[str, Any]]) -> None:
|
|
"""Notify all providers of session end."""
|
|
for provider in self._providers:
|
|
try:
|
|
provider.on_session_end(messages)
|
|
except Exception as e:
|
|
logger.debug(
|
|
"Memory provider '%s' on_session_end failed: %s",
|
|
provider.name, e,
|
|
)
|
|
|
|
def on_pre_compress(self, messages: List[Dict[str, Any]]) -> None:
|
|
"""Notify all providers before context compression."""
|
|
for provider in self._providers:
|
|
try:
|
|
provider.on_pre_compress(messages)
|
|
except Exception as e:
|
|
logger.debug(
|
|
"Memory provider '%s' on_pre_compress failed: %s",
|
|
provider.name, e,
|
|
)
|
|
|
|
def on_memory_write(self, action: str, target: str, content: str) -> None:
|
|
"""Notify external providers when the built-in memory tool writes.
|
|
|
|
Skips the builtin provider itself (it's the source of the write).
|
|
"""
|
|
for provider in self._providers:
|
|
if provider.name == "builtin":
|
|
continue
|
|
try:
|
|
provider.on_memory_write(action, target, content)
|
|
except Exception as e:
|
|
logger.debug(
|
|
"Memory provider '%s' on_memory_write failed: %s",
|
|
provider.name, e,
|
|
)
|
|
|
|
def shutdown_all(self) -> None:
|
|
"""Shut down all providers (reverse order for clean teardown)."""
|
|
for provider in reversed(self._providers):
|
|
try:
|
|
provider.shutdown()
|
|
except Exception as e:
|
|
logger.warning(
|
|
"Memory provider '%s' shutdown failed: %s",
|
|
provider.name, e,
|
|
)
|
|
|
|
def initialize_all(self, session_id: str, **kwargs) -> None:
|
|
"""Initialize all providers."""
|
|
for provider in self._providers:
|
|
try:
|
|
provider.initialize(session_id=session_id, **kwargs)
|
|
except Exception as e:
|
|
logger.warning(
|
|
"Memory provider '%s' initialize failed: %s",
|
|
provider.name, e,
|
|
)
|