0ef5fcb1c5
Security / Dependency audit (pip-audit) (push) Has been cancelled
Security / CodeQL (javascript-typescript) (push) Has been cancelled
Security / CodeQL (python) (push) Has been cancelled
Security / Secret scan (gitleaks) (push) Has been cancelled
rust / test (ubuntu) (push) Has been cancelled
rust / simulator e2e (macos-latest) (push) Has been cancelled
rust / simulator e2e (ubuntu-latest) (push) Has been cancelled
rust / simulator e2e (windows-latest) (push) Has been cancelled
rust / wheels (aarch64-apple-darwin) (push) Has been cancelled
rust / wheels (x86_64-unknown-linux-gnu) (push) Has been cancelled
rust / wheels (x86_64-apple-darwin) (push) Has been cancelled
rust / audit (push) Has been cancelled
rust / parity (nightly, allowed to fail during Phase 0) (push) Has been cancelled
CI / commitlint (push) Has been skipped
Dev Containers / validate (.devcontainer/devcontainer.json, default) (push) Failing after 0s
Dev Containers / validate (.devcontainer/memory-stack/devcontainer.json, memory-stack) (push) Failing after 0s
Dev Containers / validate-worktree (push) Failing after 0s
CI / changes (push) Failing after 4s
Deploy Documentation / validate (push) Has been skipped
Deploy Documentation / deploy (push) Failing after 1s
Init Native E2E / init-native (ubuntu-latest, claude) (push) Failing after 1s
Init Native E2E / init-native (ubuntu-latest, codex) (push) Failing after 1s
Install Native E2E / install-native (ubuntu-latest) (push) Failing after 1s
OpenCode Plugin / typecheck + build + test (push) Failing after 1s
Init Native E2E / init-native (ubuntu-latest, copilot) (push) Failing after 1s
Release Please / release-please (push) Failing after 1s
Wrap E2E / docker-wrap-e2e (push) Failing after 1s
Wrap Native E2E / wrap-native (ubuntu-latest) (push) Failing after 1s
Init E2E / docker-init-e2e (push) Failing after 4s
Merge Conflicts / merge-conflicts (push) Failing after 4s
CI / lint (push) Has been cancelled
CI / build-wheel (push) Has been cancelled
CI / build-wheel-windows (push) Has been cancelled
CI / prefetch-model (push) Has been cancelled
CI / test-dashboard-ui (push) Has been cancelled
CI / test (1) (push) Has been cancelled
CI / test (2) (push) Has been cancelled
CI / test (3) (push) Has been cancelled
CI / test (4) (push) Has been cancelled
CI / test-extras (push) Has been cancelled
CI / test-agno (push) Has been cancelled
CI / build (push) Has been cancelled
CI / workflow-validation (push) Has been cancelled
CI / docker-native-e2e (push) Has been cancelled
CI / windows-native-wrapper (push) Has been cancelled
CI / macos-native-wrapper (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code-nonroot name:code-nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code-slim name:code-slim]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code-slim-nonroot name:code-slim-nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-nonroot name:nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-slim name:slim]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-slim-nonroot name:slim-nonroot]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime name:]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-code name:code]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-code-nonroot name:code-nonroot]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-code-slim name:code-slim]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-code-slim-nonroot name:code-slim-nonroot]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-nonroot name:nonroot]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-slim name:slim]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-slim-nonroot name:slim-nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime name:]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code name:code]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code-nonroot name:code-nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code-slim name:code-slim]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code-slim-nonroot name:code-slim-nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-nonroot name:nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-slim name:slim]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-slim-nonroot name:slim-nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime name:]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code name:code]) (push) Has been cancelled
Docker / promote-latest (push) Has been cancelled
Init Native E2E / init-native (macos-latest, claude) (push) Has been cancelled
Init Native E2E / init-native (macos-latest, codex) (push) Has been cancelled
Init Native E2E / init-native (macos-latest, copilot) (push) Has been cancelled
Install Native E2E / install-native (macos-latest) (push) Has been cancelled
Wrap Native E2E / wrap-native (macos-latest) (push) Has been cancelled
894 lines
34 KiB
Python
894 lines
34 KiB
Python
"""Local backend adapter for Headroom's hierarchical memory system.
|
|
|
|
Provides a fully local memory backend using embedded databases:
|
|
- SQLite for memory storage
|
|
- SQLite-vec for vector search (bounded, persistent) - preferred
|
|
- HNSW for vector search (fallback if sqlite-vec unavailable)
|
|
- FTS5 for text search
|
|
- SQLite graph for relationships (bounded memory, persistent)
|
|
|
|
No network calls required, fast startup, suitable for development and
|
|
single-process production deployments.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from headroom.memory.adapters.graph_models import Entity, Relationship, Subgraph
|
|
from headroom.memory.models import Memory
|
|
from headroom.memory.ports import MemorySearchResult
|
|
from headroom.models.config import ML_MODEL_DEFAULTS
|
|
|
|
if TYPE_CHECKING:
|
|
from headroom.memory.adapters.graph import InMemoryGraphStore
|
|
from headroom.memory.adapters.sqlite_graph import SQLiteGraphStore
|
|
from headroom.memory.core import HierarchicalMemory
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class LocalBackendConfig:
|
|
"""Configuration for local backend.
|
|
|
|
Attributes:
|
|
db_path: Path to the SQLite database file for memories.
|
|
graph_db_path: Path to the SQLite database file for graph. If None,
|
|
derives from db_path (e.g., "memory.db" -> "memory_graph.db").
|
|
embedder_model: Name of the sentence-transformers model for embeddings.
|
|
vector_dimension: Dimension of embedding vectors (must match embedder model).
|
|
graph_persist: If True, use SQLiteGraphStore (bounded, persistent).
|
|
If False, use InMemoryGraphStore (unbounded, volatile).
|
|
graph_cache_size_kb: SQLite page cache size for graph store in KB.
|
|
Higher = more memory, faster queries. Default: 8192 (8MB).
|
|
cache_enabled: Whether to enable memory caching.
|
|
cache_max_size: Maximum number of entries in the cache.
|
|
"""
|
|
|
|
db_path: str = "memory.db"
|
|
graph_db_path: str | None = None # Derived from db_path if not specified
|
|
embedder_backend: str = "local" # "local" (sentence-transformers), "openai", "ollama"
|
|
embedder_model: str = field(default_factory=lambda: ML_MODEL_DEFAULTS.sentence_transformer)
|
|
vector_dimension: int = field(
|
|
default_factory=lambda: ML_MODEL_DEFAULTS.sentence_transformer_dim
|
|
)
|
|
openai_api_key: str | None = None # Required when embedder_backend="openai"
|
|
ollama_base_url: str = "http://localhost:11434" # For embedder_backend="ollama"
|
|
graph_persist: bool = True # Use SQLiteGraphStore (bounded, persistent)
|
|
graph_cache_size_kb: int = 8192 # 8MB default
|
|
cache_enabled: bool = True
|
|
cache_max_size: int = 1000
|
|
|
|
|
|
class LocalBackend:
|
|
"""
|
|
Local backend using embedded databases.
|
|
|
|
This backend provides a fully local memory system with:
|
|
- SQLite for memory storage (MemoryStore)
|
|
- SQLite-vec for vector search (VectorIndex) - bounded, persistent
|
|
- FTS5 for text search (TextIndex)
|
|
- SQLite graph for relationships (GraphStore) - bounded, persistent
|
|
|
|
All operations are performed locally with no network calls,
|
|
making it suitable for:
|
|
- Development and testing
|
|
- Single-process applications
|
|
- Privacy-sensitive deployments
|
|
- Offline operation
|
|
|
|
Usage:
|
|
config = LocalBackendConfig(
|
|
db_path="my_memory.db",
|
|
embedder_model="all-MiniLM-L6-v2",
|
|
)
|
|
backend = LocalBackend(config)
|
|
|
|
# Save a memory with entities and relationships
|
|
memory = await backend.save_memory(
|
|
content="Alice works at Acme Corp",
|
|
user_id="user123",
|
|
importance=0.8,
|
|
entities=["Alice", "Acme Corp"],
|
|
relationships=[{"source": "Alice", "target": "Acme Corp", "type": "works_at"}],
|
|
)
|
|
|
|
# Search with graph expansion
|
|
results = await backend.search_memories(
|
|
query="Where does Alice work?",
|
|
user_id="user123",
|
|
include_related=True,
|
|
)
|
|
"""
|
|
|
|
def __init__(self, config: LocalBackendConfig | None = None) -> None:
|
|
"""Initialize the local backend.
|
|
|
|
Args:
|
|
config: Configuration for the backend. Uses defaults if None.
|
|
"""
|
|
self._config = config or LocalBackendConfig()
|
|
self._initialized = False
|
|
self._hierarchical_memory: HierarchicalMemory | None = None
|
|
self._graph: InMemoryGraphStore | SQLiteGraphStore | None = None
|
|
# Async singleflight guard for lazy init. Per-project backends handed
|
|
# out by BackendRouter init lazily on first use; concurrent first
|
|
# callers must land on ONE init (double-checked pattern below) instead
|
|
# of racing N partial inits that leave ``_hierarchical_memory`` None
|
|
# and trip the ``assert`` guards downstream. Created lazily so the
|
|
# backend can be constructed before an event loop exists.
|
|
self._init_lock: asyncio.Lock | None = None
|
|
|
|
def _get_init_lock(self) -> asyncio.Lock:
|
|
"""Lazily create the init lock bound to the running event loop."""
|
|
if self._init_lock is None:
|
|
self._init_lock = asyncio.Lock()
|
|
return self._init_lock
|
|
|
|
async def _ensure_initialized(self) -> None:
|
|
"""Ensure the backend is initialized with all components.
|
|
|
|
Creates the HierarchicalMemory system and graph store on first use.
|
|
Uses SQLiteGraphStore (bounded, persistent) when graph_persist=True,
|
|
or InMemoryGraphStore (unbounded, volatile) when graph_persist=False.
|
|
|
|
Singleflight via ``self._init_lock`` with a double-checked flag:
|
|
concurrent first callers await the same cold-start (which can exceed
|
|
a second on the ``pytorch_mps`` embedder) rather than each kicking off
|
|
a parallel init. If a slow init is cancelled (e.g. an outer
|
|
``asyncio.wait_for`` timeout), state is reset so a later call retries
|
|
cleanly and ``CancelledError`` is re-raised rather than leaving the
|
|
backend half-built.
|
|
"""
|
|
# Fast path: already initialized, no lock contention.
|
|
if self._initialized:
|
|
return
|
|
|
|
lock = self._get_init_lock()
|
|
async with lock:
|
|
# Double-check after acquiring the lock — another task may have
|
|
# finished the init while we were waiting.
|
|
if self._initialized:
|
|
return
|
|
try:
|
|
await self._init_locked()
|
|
except asyncio.CancelledError:
|
|
# Cancellation (e.g. wait_for timeout) can leave a partial
|
|
# backend. Reset so the next call re-inits from scratch and
|
|
# never sees a half-built ``_hierarchical_memory``.
|
|
self._hierarchical_memory = None
|
|
self._graph = None
|
|
self._initialized = False
|
|
raise
|
|
|
|
async def _init_locked(self) -> None:
|
|
"""Actual init body. Must be called with ``_init_lock`` held."""
|
|
from headroom.memory import HierarchicalMemory, MemoryConfig
|
|
from headroom.memory.config import EmbedderBackend
|
|
|
|
# Map string embedder_backend to enum
|
|
embedder_backend_map = {
|
|
"local": EmbedderBackend.LOCAL,
|
|
"onnx": EmbedderBackend.ONNX,
|
|
"openai": EmbedderBackend.OPENAI,
|
|
"ollama": EmbedderBackend.OLLAMA,
|
|
}
|
|
embedder_backend = embedder_backend_map.get(
|
|
self._config.embedder_backend, EmbedderBackend.LOCAL
|
|
)
|
|
|
|
mem_config = MemoryConfig(
|
|
db_path=Path(self._config.db_path),
|
|
embedder_backend=embedder_backend,
|
|
embedder_model=self._config.embedder_model,
|
|
vector_dimension=self._config.vector_dimension,
|
|
openai_api_key=self._config.openai_api_key,
|
|
ollama_base_url=self._config.ollama_base_url,
|
|
cache_enabled=self._config.cache_enabled,
|
|
cache_max_size=self._config.cache_max_size,
|
|
)
|
|
|
|
self._hierarchical_memory = await HierarchicalMemory.create(mem_config)
|
|
|
|
# Choose graph store based on config
|
|
if self._config.graph_persist:
|
|
from headroom.memory.adapters.sqlite_graph import SQLiteGraphStore
|
|
|
|
# Derive graph db path from main db path if not specified
|
|
if self._config.graph_db_path:
|
|
graph_db_path = self._config.graph_db_path
|
|
else:
|
|
# "memory.db" -> "memory_graph.db"
|
|
db_path = Path(self._config.db_path)
|
|
graph_db_path = str(db_path.parent / f"{db_path.stem}_graph{db_path.suffix}")
|
|
|
|
self._graph = SQLiteGraphStore(
|
|
db_path=graph_db_path,
|
|
page_cache_size_kb=self._config.graph_cache_size_kb,
|
|
)
|
|
logger.info(
|
|
f"LocalBackend: Using SQLiteGraphStore at {graph_db_path} "
|
|
f"(cache: {self._config.graph_cache_size_kb}KB)"
|
|
)
|
|
else:
|
|
from headroom.memory.adapters.graph import InMemoryGraphStore
|
|
|
|
self._graph = InMemoryGraphStore()
|
|
logger.info("LocalBackend: Using InMemoryGraphStore (unbounded)")
|
|
|
|
self._initialized = True
|
|
|
|
# =========================================================================
|
|
# Core Memory Operations
|
|
# =========================================================================
|
|
|
|
async def save_memory(
|
|
self,
|
|
content: str,
|
|
user_id: str,
|
|
importance: float = 0.5,
|
|
entities: list[str] | None = None,
|
|
relationships: list[dict[str, Any]] | None = None,
|
|
metadata: dict[str, Any] | None = None,
|
|
session_id: str | None = None,
|
|
agent_id: str | None = None,
|
|
turn_id: str | None = None,
|
|
# Pre-extraction fields (for optimized mode)
|
|
facts: list[str] | None = None,
|
|
extracted_entities: list[dict[str, str]] | None = None,
|
|
extracted_relationships: list[dict[str, str]] | None = None,
|
|
) -> Memory:
|
|
"""Save a memory with optional entities and relationships.
|
|
|
|
Creates a memory, stores it via HierarchicalMemory, and optionally
|
|
adds entities and relationships to the knowledge graph.
|
|
|
|
Supports two modes:
|
|
1. Standard mode: Pass content, entities, and relationships (inferred types)
|
|
2. Pre-extraction mode: Pass facts, extracted_entities, extracted_relationships
|
|
(from optimized tool schema with explicit types)
|
|
|
|
Args:
|
|
content: The memory content/text.
|
|
user_id: User identifier (required).
|
|
importance: Importance score (0.0 - 1.0).
|
|
entities: List of entity names to add to the graph (simple format).
|
|
relationships: List of relationship dicts with keys:
|
|
- source: Source entity name
|
|
- target: Target entity name
|
|
- type: Relationship type (e.g., "works_at", "knows")
|
|
metadata: Additional metadata.
|
|
session_id: Optional session identifier.
|
|
agent_id: Optional agent identifier.
|
|
turn_id: Optional turn identifier.
|
|
facts: Pre-extracted discrete facts (optimized mode).
|
|
If provided, each fact is stored as a separate memory.
|
|
extracted_entities: Pre-extracted entities with types (optimized mode).
|
|
Format: [{"entity": "name", "entity_type": "type"}]
|
|
extracted_relationships: Pre-extracted relationships (optimized mode).
|
|
Format: [{"source": "entity1", "relationship": "type", "destination": "entity2"}]
|
|
|
|
Returns:
|
|
The created Memory object (for the main content, or first fact if facts provided).
|
|
"""
|
|
await self._ensure_initialized()
|
|
assert self._hierarchical_memory is not None
|
|
assert self._graph is not None
|
|
|
|
# Determine if using pre-extraction mode
|
|
has_pre_extraction = bool(facts or extracted_entities or extracted_relationships)
|
|
|
|
# Merge entity names from both simple and typed formats
|
|
all_entity_names: list[str] = list(entities) if entities else []
|
|
entity_types: dict[str, str] = {}
|
|
|
|
if extracted_entities:
|
|
for ent in extracted_entities:
|
|
name = ent.get("entity", "")
|
|
if name and name not in all_entity_names:
|
|
all_entity_names.append(name)
|
|
if name:
|
|
entity_types[name.lower()] = ent.get("entity_type", "unknown")
|
|
|
|
# Merge relationships from both formats
|
|
all_relationships: list[dict[str, Any]] = list(relationships) if relationships else []
|
|
|
|
if extracted_relationships:
|
|
for rel in extracted_relationships:
|
|
# Convert from optimized format to standard format
|
|
all_relationships.append(
|
|
{
|
|
"source": rel.get("source", ""),
|
|
"target": rel.get("destination", ""), # Note: "destination" in optimized
|
|
"type": rel.get("relationship", "related_to"),
|
|
}
|
|
)
|
|
|
|
# Prepare base metadata
|
|
base_metadata = metadata or {}
|
|
if has_pre_extraction:
|
|
base_metadata["_pre_extracted"] = True
|
|
if facts:
|
|
base_metadata["_fact_count"] = len(facts)
|
|
|
|
# Store memories
|
|
memories_created: list[Memory] = []
|
|
|
|
if facts:
|
|
# Store each fact as a separate memory (like DirectMem0Adapter)
|
|
for i, fact in enumerate(facts):
|
|
fact_metadata = {**base_metadata, "_fact_index": i}
|
|
memory = await self._hierarchical_memory.add(
|
|
content=fact,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
agent_id=agent_id,
|
|
turn_id=turn_id,
|
|
importance=importance,
|
|
entity_refs=all_entity_names,
|
|
metadata=fact_metadata,
|
|
)
|
|
memories_created.append(memory)
|
|
else:
|
|
# Store single memory with content
|
|
memory = await self._hierarchical_memory.add(
|
|
content=content,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
agent_id=agent_id,
|
|
turn_id=turn_id,
|
|
importance=importance,
|
|
entity_refs=all_entity_names,
|
|
metadata=base_metadata,
|
|
)
|
|
memories_created.append(memory)
|
|
|
|
# Get primary memory (first created) for graph linking
|
|
primary_memory = memories_created[0]
|
|
|
|
# Add entities to graph
|
|
if all_entity_names:
|
|
entity_id_map: dict[str, str] = {}
|
|
|
|
for entity_name in all_entity_names:
|
|
# Check if entity already exists
|
|
existing = await self._graph.get_entity_by_name(user_id, entity_name)
|
|
if existing:
|
|
entity_id_map[entity_name.lower()] = existing.id
|
|
else:
|
|
# Create new entity with type if available
|
|
entity_type = entity_types.get(entity_name.lower(), "unknown")
|
|
entity = Entity(
|
|
id=str(uuid.uuid4()),
|
|
user_id=user_id,
|
|
name=entity_name,
|
|
entity_type=entity_type,
|
|
metadata={"source_memory_id": primary_memory.id},
|
|
)
|
|
await self._graph.add_entity(entity)
|
|
entity_id_map[entity_name.lower()] = entity.id
|
|
|
|
# Add relationships to graph
|
|
if all_relationships:
|
|
for rel in all_relationships:
|
|
source_name = rel.get("source", "").lower()
|
|
target_name = rel.get("target", "").lower()
|
|
rel_type = rel.get("type", "related_to")
|
|
|
|
source_id = entity_id_map.get(source_name)
|
|
target_id = entity_id_map.get(target_name)
|
|
|
|
if source_id and target_id:
|
|
relationship = Relationship(
|
|
id=str(uuid.uuid4()),
|
|
user_id=user_id,
|
|
source_id=source_id,
|
|
target_id=target_id,
|
|
relation_type=rel_type,
|
|
metadata={"source_memory_id": primary_memory.id},
|
|
)
|
|
await self._graph.add_relationship(relationship)
|
|
|
|
# Return primary memory
|
|
return primary_memory
|
|
|
|
async def search_memories(
|
|
self,
|
|
query: str,
|
|
user_id: str,
|
|
top_k: int = 10,
|
|
entities: list[str] | None = None,
|
|
include_related: bool = True,
|
|
min_similarity: float = 0.0,
|
|
session_id: str | None = None,
|
|
) -> list[MemorySearchResult]:
|
|
"""Search memories with optional graph expansion.
|
|
|
|
Performs vector search via HierarchicalMemory and optionally
|
|
expands results via the knowledge graph.
|
|
|
|
Args:
|
|
query: Natural language search query.
|
|
user_id: User identifier to scope the search.
|
|
top_k: Maximum number of results to return.
|
|
entities: Optional filter by related entities.
|
|
include_related: If True, expand results via knowledge graph.
|
|
min_similarity: Minimum cosine similarity threshold.
|
|
session_id: Optional session filter to isolate memories by session.
|
|
|
|
Returns:
|
|
List of MemorySearchResult objects with scores and related entities.
|
|
"""
|
|
await self._ensure_initialized()
|
|
assert self._hierarchical_memory is not None
|
|
assert self._graph is not None
|
|
|
|
# Perform vector search
|
|
vector_results = await self._hierarchical_memory.search(
|
|
query=query,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
top_k=top_k * 2 if include_related else top_k, # Over-fetch for deduplication
|
|
min_similarity=min_similarity,
|
|
)
|
|
|
|
# Convert to MemorySearchResult and collect entity refs
|
|
results: list[MemorySearchResult] = []
|
|
seen_memory_ids: set[str] = set()
|
|
all_entity_refs: set[str] = set()
|
|
|
|
for vr in vector_results:
|
|
if vr.memory.id in seen_memory_ids:
|
|
continue
|
|
|
|
seen_memory_ids.add(vr.memory.id)
|
|
all_entity_refs.update(vr.memory.entity_refs)
|
|
|
|
results.append(
|
|
MemorySearchResult(
|
|
memory=vr.memory,
|
|
score=vr.similarity,
|
|
related_entities=list(vr.memory.entity_refs),
|
|
related_memories=[],
|
|
)
|
|
)
|
|
|
|
# Graph expansion if requested
|
|
if include_related and all_entity_refs:
|
|
# Find entities in graph by name
|
|
entity_ids: list[str] = []
|
|
for entity_name in all_entity_refs:
|
|
entity = await self._graph.get_entity_by_name(user_id, entity_name)
|
|
if entity:
|
|
entity_ids.append(entity.id)
|
|
|
|
if entity_ids:
|
|
# Query subgraph with 1-2 hops
|
|
subgraph = await self._graph.query_subgraph(
|
|
entity_ids=entity_ids,
|
|
max_hops=2,
|
|
)
|
|
|
|
# Get memory IDs linked to discovered entities
|
|
related_memory_ids: set[str] = set()
|
|
for entity in subgraph.entities:
|
|
if entity.metadata and "source_memory_id" in entity.metadata:
|
|
related_memory_ids.add(entity.metadata["source_memory_id"])
|
|
for rel in subgraph.relationships:
|
|
if rel.metadata and "source_memory_id" in rel.metadata:
|
|
related_memory_ids.add(rel.metadata["source_memory_id"])
|
|
|
|
# Fetch related memories not already in results
|
|
new_memory_ids = related_memory_ids - seen_memory_ids
|
|
for mem_id in new_memory_ids:
|
|
memory = await self._hierarchical_memory.get(mem_id)
|
|
if memory and memory.user_id == user_id:
|
|
# Filter by session_id if specified (security: prevent session leakage)
|
|
if session_id is not None and memory.session_id != session_id:
|
|
continue
|
|
# Add with lower score since it's from graph expansion
|
|
results.append(
|
|
MemorySearchResult(
|
|
memory=memory,
|
|
score=0.5, # Default score for graph-expanded results
|
|
related_entities=list(memory.entity_refs),
|
|
related_memories=[],
|
|
)
|
|
)
|
|
seen_memory_ids.add(mem_id)
|
|
|
|
# Filter by specified entities if provided
|
|
if entities:
|
|
entities_lower = {e.lower() for e in entities}
|
|
results = [
|
|
r
|
|
for r in results
|
|
if any(ref.lower() in entities_lower for ref in r.related_entities)
|
|
]
|
|
|
|
# Sort by score and limit
|
|
results.sort(key=lambda x: x.score, reverse=True)
|
|
return results[:top_k]
|
|
|
|
async def record_access(self, memory_ids: list[str]) -> int:
|
|
"""Record retrieval metadata for memories returned to a caller."""
|
|
await self._ensure_initialized()
|
|
assert self._hierarchical_memory is not None
|
|
return await self._hierarchical_memory.record_access(memory_ids)
|
|
|
|
async def update_memory(
|
|
self,
|
|
memory_id: str,
|
|
new_content: str,
|
|
reason: str | None = None,
|
|
user_id: str | None = None,
|
|
) -> Memory:
|
|
"""Update a memory with new content (creates versioned history).
|
|
|
|
Uses HierarchicalMemory.supersede() to create a new version while
|
|
preserving the old version for historical queries.
|
|
|
|
Args:
|
|
memory_id: ID of the memory to update.
|
|
new_content: The new content for the memory.
|
|
reason: Reason for the update (stored for audit trail).
|
|
user_id: User ID for validation (optional).
|
|
|
|
Returns:
|
|
The new Memory that supersedes the old one.
|
|
|
|
Raises:
|
|
ValueError: If the memory is not found.
|
|
"""
|
|
# Note: reason and user_id are accepted for protocol compliance
|
|
# but not yet used in the underlying implementation
|
|
_ = reason
|
|
_ = user_id
|
|
await self._ensure_initialized()
|
|
assert self._hierarchical_memory is not None
|
|
|
|
# Use supersede for versioned updates
|
|
new_memory = await self._hierarchical_memory.supersede(
|
|
old_memory_id=memory_id,
|
|
new_content=new_content,
|
|
)
|
|
|
|
return new_memory
|
|
|
|
async def delete_memory(
|
|
self,
|
|
memory_id: str,
|
|
reason: str | None = None,
|
|
user_id: str | None = None,
|
|
) -> bool:
|
|
"""Delete a memory.
|
|
|
|
Removes the memory from storage and indexes. Also cleans up
|
|
graph references if the memory was the source of entities
|
|
or relationships.
|
|
|
|
Args:
|
|
memory_id: ID of the memory to delete.
|
|
reason: Reason for deletion (stored for audit trail).
|
|
user_id: User ID for validation (optional).
|
|
|
|
Returns:
|
|
True if the memory was deleted, False if not found.
|
|
"""
|
|
# Note: reason and user_id are accepted for protocol compliance
|
|
# but not yet used in the underlying implementation
|
|
_ = reason
|
|
_ = user_id
|
|
await self._ensure_initialized()
|
|
assert self._hierarchical_memory is not None
|
|
assert self._graph is not None
|
|
|
|
# Get memory to find its user_id for graph cleanup
|
|
memory = await self._hierarchical_memory.get(memory_id)
|
|
if memory is None:
|
|
return False
|
|
|
|
# Delete from HierarchicalMemory
|
|
deleted = await self._hierarchical_memory.delete(memory_id)
|
|
|
|
if deleted:
|
|
# Clean up graph references
|
|
# Find and delete entities that were created from this memory
|
|
user_entities = await self._graph.get_entities_for_user(memory.user_id)
|
|
for entity in user_entities:
|
|
if entity.metadata.get("source_memory_id") == memory_id:
|
|
await self._graph.delete_entity(entity.id)
|
|
|
|
return deleted
|
|
|
|
async def get_memory(self, memory_id: str) -> Memory | None:
|
|
"""Get a memory by ID.
|
|
|
|
Args:
|
|
memory_id: The unique identifier of the memory.
|
|
|
|
Returns:
|
|
The Memory if found, None otherwise.
|
|
"""
|
|
await self._ensure_initialized()
|
|
assert self._hierarchical_memory is not None
|
|
|
|
return await self._hierarchical_memory.get(memory_id)
|
|
|
|
# =========================================================================
|
|
# Capability Properties
|
|
# =========================================================================
|
|
|
|
@property
|
|
def supports_graph(self) -> bool:
|
|
"""Whether this backend supports knowledge graph operations.
|
|
|
|
Returns:
|
|
True, as this backend uses SQLiteGraphStore (default) or InMemoryGraphStore.
|
|
"""
|
|
return True
|
|
|
|
@property
|
|
def supports_vector_search(self) -> bool:
|
|
"""Whether this backend supports semantic vector search.
|
|
|
|
Returns:
|
|
True, as this backend uses HNSW vector index.
|
|
"""
|
|
return True
|
|
|
|
@property
|
|
def supports_text_search(self) -> bool:
|
|
"""Whether this backend supports full-text search.
|
|
|
|
Returns:
|
|
True, as this backend uses FTS5 text index.
|
|
"""
|
|
return True
|
|
|
|
# =========================================================================
|
|
# Graph Operations
|
|
# =========================================================================
|
|
|
|
async def get_graph(self) -> InMemoryGraphStore | SQLiteGraphStore:
|
|
"""Get the underlying graph store.
|
|
|
|
Returns:
|
|
The graph store instance (SQLiteGraphStore if graph_persist=True,
|
|
InMemoryGraphStore otherwise).
|
|
"""
|
|
await self._ensure_initialized()
|
|
assert self._graph is not None
|
|
return self._graph
|
|
|
|
async def query_subgraph(
|
|
self,
|
|
entity_names: list[str],
|
|
user_id: str,
|
|
max_hops: int = 2,
|
|
) -> Subgraph:
|
|
"""Query a subgraph starting from named entities.
|
|
|
|
Args:
|
|
entity_names: Starting entity names for the traversal.
|
|
user_id: User identifier for entity lookup.
|
|
max_hops: Maximum number of hops from starting entities.
|
|
|
|
Returns:
|
|
Subgraph containing reachable entities and relationships.
|
|
"""
|
|
await self._ensure_initialized()
|
|
assert self._graph is not None
|
|
|
|
# Find entity IDs by name
|
|
entity_ids: list[str] = []
|
|
for name in entity_names:
|
|
entity = await self._graph.get_entity_by_name(user_id, name)
|
|
if entity:
|
|
entity_ids.append(entity.id)
|
|
|
|
if not entity_ids:
|
|
return Subgraph(entities=[], relationships=[], root_entity_ids=[])
|
|
|
|
return await self._graph.query_subgraph(
|
|
entity_ids=entity_ids,
|
|
max_hops=max_hops,
|
|
)
|
|
|
|
# =========================================================================
|
|
# Additional Convenience Methods
|
|
# =========================================================================
|
|
|
|
async def get_user_memories(
|
|
self,
|
|
user_id: str,
|
|
limit: int = 100,
|
|
) -> list[Memory]:
|
|
"""Get all memories for a user.
|
|
|
|
Args:
|
|
user_id: User identifier.
|
|
limit: Maximum number of memories to return.
|
|
|
|
Returns:
|
|
List of memories for the user.
|
|
"""
|
|
await self._ensure_initialized()
|
|
assert self._hierarchical_memory is not None
|
|
|
|
return await self._hierarchical_memory.get_user_memories(
|
|
user_id=user_id,
|
|
limit=limit,
|
|
)
|
|
|
|
async def clear_user(self, user_id: str) -> int:
|
|
"""Clear all memories and graph data for a user.
|
|
|
|
Args:
|
|
user_id: User identifier.
|
|
|
|
Returns:
|
|
Number of memories deleted.
|
|
"""
|
|
await self._ensure_initialized()
|
|
assert self._hierarchical_memory is not None
|
|
assert self._graph is not None
|
|
|
|
# Clear memories
|
|
count = await self._hierarchical_memory.clear_scope(user_id=user_id)
|
|
|
|
# Clear graph
|
|
await self._graph.clear_user(user_id)
|
|
|
|
return count
|
|
|
|
async def close(self) -> None:
|
|
"""Close the backend and release resources."""
|
|
# Close HierarchicalMemory to release httpx clients in embedders
|
|
if self._hierarchical_memory is not None:
|
|
await self._hierarchical_memory.close()
|
|
self._hierarchical_memory = None
|
|
self._graph = None
|
|
self._initialized = False
|
|
|
|
# =========================================================================
|
|
# Text Search
|
|
# =========================================================================
|
|
|
|
async def text_search(
|
|
self,
|
|
query: str,
|
|
user_id: str,
|
|
limit: int = 100,
|
|
) -> list[MemorySearchResult]:
|
|
"""Full-text search for memories.
|
|
|
|
Uses FTS5 index for keyword matching with BM25 ranking.
|
|
|
|
Args:
|
|
query: Search query text.
|
|
user_id: User identifier to scope the search.
|
|
limit: Maximum number of results.
|
|
|
|
Returns:
|
|
List of MemorySearchResult objects.
|
|
"""
|
|
await self._ensure_initialized()
|
|
assert self._hierarchical_memory is not None
|
|
|
|
# Build text filter and call text index directly
|
|
# (HierarchicalMemory.text_search has a signature mismatch with the protocol)
|
|
from headroom.memory.ports import TextFilter
|
|
|
|
text_filter = TextFilter(
|
|
query=query,
|
|
user_id=user_id,
|
|
limit=limit,
|
|
)
|
|
|
|
# Use the protocol-compliant search_memories method on the text index
|
|
text_index = self._hierarchical_memory.text_index
|
|
text_results = await text_index.search_memories(text_filter) # type: ignore[attr-defined]
|
|
|
|
# Convert to MemorySearchResult
|
|
return [
|
|
MemorySearchResult(
|
|
memory=tr.memory,
|
|
score=tr.score,
|
|
related_entities=list(tr.memory.entity_refs),
|
|
related_memories=[],
|
|
)
|
|
for tr in text_results
|
|
]
|
|
|
|
async def hybrid_search(
|
|
self,
|
|
query: str,
|
|
user_id: str,
|
|
top_k: int = 10,
|
|
vector_weight: float = 0.5,
|
|
text_weight: float = 0.5,
|
|
min_similarity: float = 0.0,
|
|
) -> list[MemorySearchResult]:
|
|
"""Hybrid search combining vector similarity and text matching.
|
|
|
|
Performs both semantic (vector) and keyword (BM25) search,
|
|
then merges results with weighted score combination.
|
|
|
|
Args:
|
|
query: Search query text.
|
|
user_id: User identifier to scope the search.
|
|
top_k: Maximum number of results to return.
|
|
vector_weight: Weight for vector similarity scores (0-1).
|
|
text_weight: Weight for text match scores (0-1).
|
|
min_similarity: Minimum similarity threshold for vector results.
|
|
|
|
Returns:
|
|
List of MemorySearchResult objects sorted by combined score.
|
|
"""
|
|
await self._ensure_initialized()
|
|
|
|
# Fetch more candidates than needed for better coverage
|
|
fetch_k = top_k * 3
|
|
|
|
# Perform vector search
|
|
vector_results = await self.search_memories(
|
|
query=query,
|
|
user_id=user_id,
|
|
top_k=fetch_k,
|
|
include_related=False,
|
|
min_similarity=min_similarity,
|
|
)
|
|
|
|
# Perform text search
|
|
text_results = await self.text_search(
|
|
query=query,
|
|
user_id=user_id,
|
|
limit=fetch_k,
|
|
)
|
|
|
|
# Normalize scores and merge
|
|
# Vector scores are already 0-1 (cosine similarity)
|
|
# Text scores need normalization
|
|
max_text_score = max((r.score for r in text_results), default=1.0) or 1.0
|
|
|
|
# Build score maps
|
|
vector_scores: dict[str, float] = {r.memory.id: r.score for r in vector_results}
|
|
text_scores: dict[str, float] = {
|
|
r.memory.id: r.score / max_text_score for r in text_results
|
|
}
|
|
|
|
# Merge all unique memories
|
|
all_memories: dict[str, MemorySearchResult] = {}
|
|
for r in vector_results:
|
|
all_memories[r.memory.id] = r
|
|
for r in text_results:
|
|
if r.memory.id not in all_memories:
|
|
all_memories[r.memory.id] = r
|
|
|
|
# Calculate combined scores
|
|
combined_results: list[tuple[float, MemorySearchResult]] = []
|
|
for memory_id, result in all_memories.items():
|
|
v_score = vector_scores.get(memory_id, 0.0)
|
|
t_score = text_scores.get(memory_id, 0.0)
|
|
combined = vector_weight * v_score + text_weight * t_score
|
|
combined_results.append((combined, result))
|
|
|
|
# Sort by combined score and return top_k
|
|
combined_results.sort(key=lambda x: x[0], reverse=True)
|
|
return [
|
|
MemorySearchResult(
|
|
memory=r.memory,
|
|
score=score,
|
|
related_entities=r.related_entities,
|
|
related_memories=r.related_memories,
|
|
)
|
|
for score, r in combined_results[:top_k]
|
|
]
|