mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
chore(memory): add EXTERNAL backend extension points
Mirrors the pattern used by headroom.ccr_backend so memory store,
vector index, and text index backends can be registered via setuptools
entry points.
- EXTERNAL enum value on StoreBackend, VectorBackend, TextBackend
- Optional *_backend_name fields on MemoryConfig
- entry_points(group=...) lookup in _create_{store,vector_index,text_index}
- New test_factory_external.py (7 tests) covering load / missing-name /
unknown-name paths
Default behavior (SQLITE + AUTO + FTS5) unchanged.
Extension groups:
headroom.memory_store
headroom.memory_vector
headroom.memory_text
This commit is contained in:
parent
e6cdc2f143
commit
5391761fe6
3 changed files with 189 additions and 3 deletions
|
|
@ -22,7 +22,7 @@ class StoreBackend(Enum):
|
|||
"""Supported memory store backends."""
|
||||
|
||||
SQLITE = "sqlite"
|
||||
# Future: POSTGRES = "postgres", DYNAMODB = "dynamodb"
|
||||
EXTERNAL = "external" # Loaded from entry_points(group="headroom.memory_store")
|
||||
|
||||
|
||||
class VectorBackend(Enum):
|
||||
|
|
@ -31,13 +31,14 @@ class VectorBackend(Enum):
|
|||
AUTO = "auto" # Auto-select: SQLITE_VEC if available, else HNSW
|
||||
SQLITE_VEC = "sqlite_vec" # SQLite-based, bounded memory, recommended
|
||||
HNSW = "hnsw" # hnswlib-based, unbounded unless max_entries set
|
||||
EXTERNAL = "external" # Loaded from entry_points(group="headroom.memory_vector")
|
||||
|
||||
|
||||
class TextBackend(Enum):
|
||||
"""Supported text index backends."""
|
||||
|
||||
FTS5 = "fts5"
|
||||
# Future: ELASTICSEARCH = "elasticsearch"
|
||||
EXTERNAL = "external" # Loaded from entry_points(group="headroom.memory_text")
|
||||
|
||||
|
||||
class EmbedderBackend(Enum):
|
||||
|
|
@ -96,10 +97,12 @@ class MemoryConfig:
|
|||
|
||||
# Storage
|
||||
store_backend: StoreBackend = StoreBackend.SQLITE
|
||||
store_backend_name: str | None = None # Required when store_backend == EXTERNAL
|
||||
db_path: Path = field(default_factory=lambda: Path("headroom_memory.db"))
|
||||
|
||||
# Vector index
|
||||
vector_backend: VectorBackend = VectorBackend.AUTO # Auto-select best available
|
||||
vector_backend_name: str | None = None # Required when vector_backend == EXTERNAL
|
||||
vector_dimension: int = 384
|
||||
vector_db_path: Path | None = (
|
||||
None # For SQLite-based vector index (derived from db_path if None)
|
||||
|
|
@ -112,6 +115,7 @@ class MemoryConfig:
|
|||
|
||||
# Text index
|
||||
text_backend: TextBackend = TextBackend.FTS5
|
||||
text_backend_name: str | None = None # Required when text_backend == EXTERNAL
|
||||
|
||||
# Embedder
|
||||
embedder_backend: EmbedderBackend = EmbedderBackend.LOCAL
|
||||
|
|
|
|||
|
|
@ -7,8 +7,9 @@ and proper wiring between components.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from importlib.metadata import entry_points
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from headroom.memory.config import (
|
||||
EmbedderBackend,
|
||||
|
|
@ -22,6 +23,37 @@ if TYPE_CHECKING:
|
|||
from headroom.memory.ports import Embedder, MemoryCache, MemoryStore, TextIndex, VectorIndex
|
||||
|
||||
|
||||
# Extension groups for memory backends registered via setuptools entry points.
|
||||
_MEMORY_STORE_GROUP = "headroom.memory_store"
|
||||
_MEMORY_VECTOR_GROUP = "headroom.memory_vector"
|
||||
_MEMORY_TEXT_GROUP = "headroom.memory_text"
|
||||
|
||||
|
||||
def _load_external_backend(
|
||||
group: str,
|
||||
name: str | None,
|
||||
field_name: str,
|
||||
config: MemoryConfig,
|
||||
) -> Any:
|
||||
"""Load a memory backend registered via setuptools entry points.
|
||||
|
||||
Mirrors the pattern used by
|
||||
`headroom.cache.compression_store._create_default_ccr_backend`.
|
||||
"""
|
||||
if not name:
|
||||
raise ValueError(
|
||||
f"{field_name} is required when backend is EXTERNAL; "
|
||||
f"set it to the entry-point name registered under '{group}'."
|
||||
)
|
||||
ep = next((e for e in entry_points(group=group) if e.name == name), None)
|
||||
if ep is None:
|
||||
raise ValueError(
|
||||
f"No entry point registered under '{group}' with name '{name}'. "
|
||||
f"Install the package that provides it."
|
||||
)
|
||||
return ep.load()(config)
|
||||
|
||||
|
||||
async def create_memory_system(
|
||||
config: MemoryConfig | None = None,
|
||||
) -> tuple[MemoryStore, VectorIndex, TextIndex, Embedder, MemoryCache | None]:
|
||||
|
|
@ -89,6 +121,14 @@ def _create_store(config: MemoryConfig) -> MemoryStore:
|
|||
|
||||
return SQLiteMemoryStore(config.db_path)
|
||||
|
||||
if config.store_backend == StoreBackend.EXTERNAL:
|
||||
return _load_external_backend( # type: ignore[no-any-return]
|
||||
_MEMORY_STORE_GROUP,
|
||||
config.store_backend_name,
|
||||
"store_backend_name",
|
||||
config,
|
||||
)
|
||||
|
||||
raise ValueError(f"Unknown store backend: {config.store_backend}")
|
||||
|
||||
|
||||
|
|
@ -149,6 +189,14 @@ def _create_vector_index(config: MemoryConfig) -> VectorIndex:
|
|||
"""
|
||||
backend = config.vector_backend
|
||||
|
||||
if backend == VectorBackend.EXTERNAL:
|
||||
return _load_external_backend( # type: ignore[no-any-return]
|
||||
_MEMORY_VECTOR_GROUP,
|
||||
config.vector_backend_name,
|
||||
"vector_backend_name",
|
||||
config,
|
||||
)
|
||||
|
||||
# AUTO: prefer SQLITE_VEC → HNSW → fail with helpful message
|
||||
if backend == VectorBackend.AUTO:
|
||||
from headroom.memory.adapters import HNSW_AVAILABLE, SQLITE_VEC_AVAILABLE
|
||||
|
|
@ -238,6 +286,14 @@ def _create_text_index(config: MemoryConfig) -> TextIndex:
|
|||
# FTS5TextIndex has a compatible interface but different method signatures
|
||||
return FTS5TextIndex(db_path=config.db_path) # type: ignore[return-value]
|
||||
|
||||
if config.text_backend == TextBackend.EXTERNAL:
|
||||
return _load_external_backend( # type: ignore[no-any-return]
|
||||
_MEMORY_TEXT_GROUP,
|
||||
config.text_backend_name,
|
||||
"text_backend_name",
|
||||
config,
|
||||
)
|
||||
|
||||
raise ValueError(f"Unknown text backend: {config.text_backend}")
|
||||
|
||||
|
||||
|
|
|
|||
126
tests/test_memory/test_factory_external.py
Normal file
126
tests/test_memory/test_factory_external.py
Normal file
|
|
@ -0,0 +1,126 @@
|
|||
"""Tests for EXTERNAL memory backends (entry-point plugins).
|
||||
|
||||
Three extension groups let packages register memory backends via
|
||||
setuptools entry points:
|
||||
- headroom.memory_store
|
||||
- headroom.memory_vector
|
||||
- headroom.memory_text
|
||||
|
||||
A package registers a callable under one of these groups; the factory
|
||||
loads it when the corresponding backend enum is EXTERNAL.
|
||||
|
||||
These tests do not require hnswlib and must remain independent of it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.memory.config import (
|
||||
MemoryConfig,
|
||||
StoreBackend,
|
||||
TextBackend,
|
||||
VectorBackend,
|
||||
)
|
||||
from headroom.memory.factory import (
|
||||
_create_store,
|
||||
_create_text_index,
|
||||
_create_vector_index,
|
||||
)
|
||||
|
||||
|
||||
class _FakeEntryPoint:
|
||||
"""Minimal stand-in for importlib.metadata.EntryPoint used in tests."""
|
||||
|
||||
def __init__(self, name: str, target):
|
||||
self.name = name
|
||||
self._target = target
|
||||
|
||||
def load(self):
|
||||
return self._target
|
||||
|
||||
|
||||
def _patch_entry_points(monkeypatch, expected_group: str, name: str, target):
|
||||
"""Patch headroom.memory.factory.entry_points to return our fake EP."""
|
||||
from headroom.memory import factory as factory_mod
|
||||
|
||||
def fake_entry_points(*, group: str):
|
||||
if group == expected_group:
|
||||
return [_FakeEntryPoint(name, target)]
|
||||
return []
|
||||
|
||||
monkeypatch.setattr(factory_mod, "entry_points", fake_entry_points)
|
||||
|
||||
|
||||
class TestExternalStoreBackend:
|
||||
"""EXTERNAL store backend loads via entry_points(group='headroom.memory_store')."""
|
||||
|
||||
def test_loads_external_store(self, monkeypatch):
|
||||
sentinel = object()
|
||||
|
||||
def make_store(config):
|
||||
assert isinstance(config, MemoryConfig)
|
||||
return sentinel
|
||||
|
||||
_patch_entry_points(monkeypatch, "headroom.memory_store", "myvec", make_store)
|
||||
|
||||
config = MemoryConfig(
|
||||
store_backend=StoreBackend.EXTERNAL,
|
||||
store_backend_name="myvec",
|
||||
)
|
||||
assert _create_store(config) is sentinel
|
||||
|
||||
def test_external_without_name_raises(self):
|
||||
config = MemoryConfig(store_backend=StoreBackend.EXTERNAL)
|
||||
with pytest.raises(ValueError, match="store_backend_name is required"):
|
||||
_create_store(config)
|
||||
|
||||
def test_external_unknown_name_raises(self, monkeypatch):
|
||||
from headroom.memory import factory as factory_mod
|
||||
|
||||
monkeypatch.setattr(factory_mod, "entry_points", lambda *, group: [])
|
||||
|
||||
config = MemoryConfig(
|
||||
store_backend=StoreBackend.EXTERNAL,
|
||||
store_backend_name="nonexistent",
|
||||
)
|
||||
with pytest.raises(ValueError, match="No entry point .* 'nonexistent'"):
|
||||
_create_store(config)
|
||||
|
||||
|
||||
class TestExternalVectorBackend:
|
||||
"""EXTERNAL vector backend loads via entry_points(group='headroom.memory_vector')."""
|
||||
|
||||
def test_loads_external_vector(self, monkeypatch):
|
||||
sentinel = object()
|
||||
_patch_entry_points(monkeypatch, "headroom.memory_vector", "myvec", lambda cfg: sentinel)
|
||||
|
||||
config = MemoryConfig(
|
||||
vector_backend=VectorBackend.EXTERNAL,
|
||||
vector_backend_name="myvec",
|
||||
)
|
||||
assert _create_vector_index(config) is sentinel
|
||||
|
||||
def test_external_without_name_raises(self):
|
||||
config = MemoryConfig(vector_backend=VectorBackend.EXTERNAL)
|
||||
with pytest.raises(ValueError, match="vector_backend_name is required"):
|
||||
_create_vector_index(config)
|
||||
|
||||
|
||||
class TestExternalTextBackend:
|
||||
"""EXTERNAL text backend loads via entry_points(group='headroom.memory_text')."""
|
||||
|
||||
def test_loads_external_text(self, monkeypatch):
|
||||
sentinel = object()
|
||||
_patch_entry_points(monkeypatch, "headroom.memory_text", "mytext", lambda cfg: sentinel)
|
||||
|
||||
config = MemoryConfig(
|
||||
text_backend=TextBackend.EXTERNAL,
|
||||
text_backend_name="mytext",
|
||||
)
|
||||
assert _create_text_index(config) is sentinel
|
||||
|
||||
def test_external_without_name_raises(self):
|
||||
config = MemoryConfig(text_backend=TextBackend.EXTERNAL)
|
||||
with pytest.raises(ValueError, match="text_backend_name is required"):
|
||||
_create_text_index(config)
|
||||
Loading…
Add table
Add a link
Reference in a new issue