Add SQLiteGraphStore for bounded, persistent graph storage

Replace unbounded InMemoryGraphStore with SQLite-backed implementation:
- Persistent storage survives proxy restarts
- Memory bounded by configurable SQLite page cache (default 8MB)
- Same async interface as InMemoryGraphStore (drop-in replacement)
- LocalBackend now uses SQLiteGraphStore by default (graph_persist=True)

New files:
- headroom/memory/adapters/sqlite_graph.py: SQLite graph store implementation
- tests/test_sqlite_graph_store.py: 37 comprehensive tests

Key features:
- O(log n) lookups via database indexes
- Case-insensitive entity name lookup per user
- BFS subgraph traversal and shortest path finding
- CASCADE delete for entity relationships
- MemoryTracker integration via get_memory_stats()
This commit is contained in:
chopratejas 2026-02-01 20:48:57 -08:00
parent e16691dd38
commit f2ecf1d2d0
4 changed files with 1715 additions and 12 deletions

View file

@ -6,6 +6,7 @@ This module provides concrete implementations of the memory system's ports:
- HNSWVectorIndex: HNSW-based vector index using hnswlib (optional)
- LRUMemoryCache: Thread-safe LRU cache for hot memories
- InMemoryGraphStore: In-memory graph store for knowledge graphs
- SQLiteGraphStore: SQLite-based graph store (bounded memory, persistent)
- LocalEmbedder: sentence-transformers embedding (local, optional)
- OpenAIEmbedder: OpenAI API embedding (cloud, optional)
- OllamaEmbedder: Ollama API embedding (local server, optional)
@ -19,6 +20,7 @@ from headroom.memory.adapters.cache import LRUMemoryCache
from headroom.memory.adapters.fts5 import FTS5TextIndex
from headroom.memory.adapters.graph import InMemoryGraphStore
from headroom.memory.adapters.sqlite import SQLiteMemoryStore
from headroom.memory.adapters.sqlite_graph import SQLiteGraphStore
# Check for optional dependencies availability
# Note: We don't import from hnsw.py here because hnswlib may crash with
@ -83,6 +85,7 @@ __all__ = [
"FTS5TextIndex",
"InMemoryGraphStore",
"LRUMemoryCache",
"SQLiteGraphStore",
"SQLiteMemoryStore",
# Optional adapters (lazy-loaded)
"HNSWVectorIndex",

View file

@ -0,0 +1,758 @@
"""SQLite graph store for Headroom's knowledge graph memory system.
Provides persistent storage for entities and relationships with efficient
lookup via database indexes and BFS-based traversal. Memory usage is bounded
by SQLite's page cache.
This is a drop-in replacement for InMemoryGraphStore that:
- Persists all data to disk
- Keeps memory bounded (configurable page cache)
- Maintains the same async interface
- Supports all query patterns (by ID, by name, BFS traversal)
"""
from __future__ import annotations
import json
import sqlite3
from collections import deque
from datetime import datetime
from pathlib import Path
from threading import RLock
from typing import TYPE_CHECKING, Any
from .graph_models import Entity, Relationship, RelationshipDirection, Subgraph
if TYPE_CHECKING:
from ..tracker import ComponentStats
class SQLiteGraphStore:
"""SQLite-based graph store implementing the GraphStore protocol.
Provides persistent storage for entities and relationships with efficient
lookup via database indexes. All operations are thread-safe.
Features:
- O(log n) entity and relationship lookup by ID (indexed)
- O(log n) entity lookup by name (per user, case-insensitive)
- O(log n) lookup of relationships by source or target entity
- BFS-based subgraph traversal with configurable hop limit
- BFS-based shortest path finding between entities
- Configurable page cache for memory bounding
Schema:
- entities: id, user_id, name, name_lower, entity_type, description,
properties, created_at, updated_at, metadata
- relationships: id, user_id, source_id, target_id, relation_type,
weight, properties, created_at, metadata
Usage:
store = SQLiteGraphStore("./graph.db")
await store.add_entity(Entity(user_id="alice", name="Project X", entity_type="project"))
entity = await store.get_entity_by_name("alice", "project x") # Case-insensitive
subgraph = await store.query_subgraph(["entity-id"], max_hops=2)
"""
def __init__(
self,
db_path: str | Path = "headroom_graph.db",
page_cache_size_kb: int = 8192, # 8MB default cache
) -> None:
"""Initialize the SQLite graph store.
Args:
db_path: Path to SQLite database file. Created if it doesn't exist.
page_cache_size_kb: SQLite page cache size in KB. Higher = more memory,
faster queries. Set to -1 for default SQLite behavior.
"""
self.db_path = Path(db_path)
self._page_cache_size_kb = page_cache_size_kb
self._lock = RLock()
self._init_db()
def _get_conn(self) -> sqlite3.Connection:
"""Get a new database connection (thread-safe pattern).
Returns:
A new SQLite connection with row factory configured.
"""
conn = sqlite3.connect(str(self.db_path))
conn.row_factory = sqlite3.Row
# Configure page cache size (negative = KB, positive = pages)
if self._page_cache_size_kb > 0:
conn.execute(f"PRAGMA cache_size = -{self._page_cache_size_kb}")
# Enable foreign keys
conn.execute("PRAGMA foreign_keys = ON")
return conn
def _init_db(self) -> None:
"""Initialize the database schema with indexes."""
with self._get_conn() as conn:
# Create entities table
conn.execute("""
CREATE TABLE IF NOT EXISTS entities (
id TEXT PRIMARY KEY,
user_id TEXT NOT NULL,
name TEXT NOT NULL,
name_lower TEXT NOT NULL,
entity_type TEXT NOT NULL DEFAULT 'unknown',
description TEXT,
properties TEXT NOT NULL DEFAULT '{}',
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL,
metadata TEXT NOT NULL DEFAULT '{}'
)
""")
# Create relationships table
conn.execute("""
CREATE TABLE IF NOT EXISTS relationships (
id TEXT PRIMARY KEY,
user_id TEXT NOT NULL,
source_id TEXT NOT NULL,
target_id TEXT NOT NULL,
relation_type TEXT NOT NULL DEFAULT 'related_to',
weight REAL NOT NULL DEFAULT 1.0,
properties TEXT NOT NULL DEFAULT '{}',
created_at TEXT NOT NULL,
metadata TEXT NOT NULL DEFAULT '{}',
FOREIGN KEY (source_id) REFERENCES entities(id) ON DELETE CASCADE,
FOREIGN KEY (target_id) REFERENCES entities(id) ON DELETE CASCADE
)
""")
# Create indexes for efficient querying
# Entity indexes
conn.execute("CREATE INDEX IF NOT EXISTS idx_entities_user_id ON entities(user_id)")
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_entities_name_lookup "
"ON entities(user_id, name_lower)"
)
conn.execute("CREATE INDEX IF NOT EXISTS idx_entities_type ON entities(entity_type)")
# Relationship indexes
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_relationships_source ON relationships(source_id)"
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_relationships_target ON relationships(target_id)"
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_relationships_type ON relationships(relation_type)"
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_relationships_user ON relationships(user_id)"
)
conn.commit()
def _entity_to_row(self, entity: Entity) -> dict[str, Any]:
"""Convert Entity object to row dict for insertion."""
return {
"id": entity.id,
"user_id": entity.user_id,
"name": entity.name,
"name_lower": entity.name.lower(),
"entity_type": entity.entity_type,
"description": entity.description,
"properties": json.dumps(entity.properties),
"created_at": entity.created_at.isoformat(),
"updated_at": entity.updated_at.isoformat(),
"metadata": json.dumps(entity.metadata),
}
def _row_to_entity(self, row: sqlite3.Row) -> Entity:
"""Convert database row to Entity object."""
return Entity(
id=row["id"],
user_id=row["user_id"],
name=row["name"],
entity_type=row["entity_type"],
description=row["description"],
properties=json.loads(row["properties"]),
created_at=datetime.fromisoformat(row["created_at"]),
updated_at=datetime.fromisoformat(row["updated_at"]),
metadata=json.loads(row["metadata"]),
)
def _relationship_to_row(self, relationship: Relationship) -> dict[str, Any]:
"""Convert Relationship object to row dict for insertion."""
return {
"id": relationship.id,
"user_id": relationship.user_id,
"source_id": relationship.source_id,
"target_id": relationship.target_id,
"relation_type": relationship.relation_type,
"weight": relationship.weight,
"properties": json.dumps(relationship.properties),
"created_at": relationship.created_at.isoformat(),
"metadata": json.dumps(relationship.metadata),
}
def _row_to_relationship(self, row: sqlite3.Row) -> Relationship:
"""Convert database row to Relationship object."""
return Relationship(
id=row["id"],
user_id=row["user_id"],
source_id=row["source_id"],
target_id=row["target_id"],
relation_type=row["relation_type"],
weight=row["weight"],
properties=json.loads(row["properties"]),
created_at=datetime.fromisoformat(row["created_at"]),
metadata=json.loads(row["metadata"]),
)
# =========================================================================
# Entity Operations
# =========================================================================
async def add_entity(self, entity: Entity) -> None:
"""Add an entity to the graph store.
If an entity with the same ID already exists, it will be replaced.
Args:
entity: The entity to add.
"""
row = self._entity_to_row(entity)
with self._lock:
with self._get_conn() as conn:
conn.execute(
"""
INSERT OR REPLACE INTO entities (
id, user_id, name, name_lower, entity_type, description,
properties, created_at, updated_at, metadata
) VALUES (
:id, :user_id, :name, :name_lower, :entity_type, :description,
:properties, :created_at, :updated_at, :metadata
)
""",
row,
)
conn.commit()
async def get_entity(self, entity_id: str) -> Entity | None:
"""Retrieve an entity by ID.
Args:
entity_id: The unique identifier of the entity.
Returns:
The entity if found, None otherwise.
"""
with self._lock:
with self._get_conn() as conn:
cursor = conn.execute(
"SELECT * FROM entities WHERE id = ?",
(entity_id,),
)
row = cursor.fetchone()
if row is None:
return None
return self._row_to_entity(row)
async def get_entity_by_name(self, user_id: str, name: str) -> Entity | None:
"""Retrieve an entity by name (case-insensitive).
Args:
user_id: The user who owns the entity.
name: The name of the entity (case-insensitive).
Returns:
The entity if found, None otherwise.
"""
with self._lock:
with self._get_conn() as conn:
cursor = conn.execute(
"SELECT * FROM entities WHERE user_id = ? AND name_lower = ?",
(user_id, name.lower()),
)
row = cursor.fetchone()
if row is None:
return None
return self._row_to_entity(row)
async def delete_entity(self, entity_id: str) -> bool:
"""Delete an entity and all its relationships.
Relationships are automatically deleted via ON DELETE CASCADE.
Args:
entity_id: The unique identifier of the entity.
Returns:
True if the entity was deleted, False if not found.
"""
with self._lock:
with self._get_conn() as conn:
cursor = conn.execute(
"DELETE FROM entities WHERE id = ?",
(entity_id,),
)
conn.commit()
return cursor.rowcount > 0
# =========================================================================
# Relationship Operations
# =========================================================================
async def add_relationship(self, relationship: Relationship) -> None:
"""Add a relationship to the graph store.
If a relationship with the same ID already exists, it will be replaced.
Args:
relationship: The relationship to add.
"""
row = self._relationship_to_row(relationship)
with self._lock:
with self._get_conn() as conn:
conn.execute(
"""
INSERT OR REPLACE INTO relationships (
id, user_id, source_id, target_id, relation_type,
weight, properties, created_at, metadata
) VALUES (
:id, :user_id, :source_id, :target_id, :relation_type,
:weight, :properties, :created_at, :metadata
)
""",
row,
)
conn.commit()
async def get_relationships(
self,
entity_id: str,
direction: RelationshipDirection = RelationshipDirection.BOTH,
relation_type: str | None = None,
) -> list[Relationship]:
"""Get relationships for an entity.
Args:
entity_id: The entity ID to get relationships for.
direction: Whether to get outgoing, incoming, or both relationships.
relation_type: Optional filter for relationship type.
Returns:
List of matching relationships.
"""
with self._lock:
with self._get_conn() as conn:
conditions = []
params: list[Any] = []
if direction == RelationshipDirection.OUTGOING:
conditions.append("source_id = ?")
params.append(entity_id)
elif direction == RelationshipDirection.INCOMING:
conditions.append("target_id = ?")
params.append(entity_id)
else: # BOTH
conditions.append("(source_id = ? OR target_id = ?)")
params.extend([entity_id, entity_id])
if relation_type is not None:
conditions.append("relation_type = ?")
params.append(relation_type)
where_clause = " AND ".join(conditions)
cursor = conn.execute(
f"SELECT * FROM relationships WHERE {where_clause}",
params,
)
return [self._row_to_relationship(row) for row in cursor]
async def delete_relationship(self, relationship_id: str) -> bool:
"""Delete a single relationship.
Args:
relationship_id: The unique identifier of the relationship.
Returns:
True if the relationship was deleted, False if not found.
"""
with self._lock:
with self._get_conn() as conn:
cursor = conn.execute(
"DELETE FROM relationships WHERE id = ?",
(relationship_id,),
)
conn.commit()
return cursor.rowcount > 0
# =========================================================================
# Graph Traversal Operations
# =========================================================================
async def query_subgraph(
self,
entity_ids: list[str],
max_hops: int = 2,
direction: RelationshipDirection = RelationshipDirection.BOTH,
relation_types: list[str] | None = None,
) -> Subgraph:
"""Query a subgraph starting from given entities using BFS traversal.
Performs a breadth-first traversal from the starting entities,
collecting all entities and relationships within the specified
number of hops.
Args:
entity_ids: Starting entity IDs for the traversal.
max_hops: Maximum number of hops from starting entities (default 2).
direction: Direction of relationship traversal.
relation_types: Optional filter for relationship types.
Returns:
Subgraph containing all reachable entities and relationships.
"""
with self._lock:
with self._get_conn() as conn:
collected_entities: dict[str, Entity] = {}
collected_relationships: dict[str, Relationship] = {}
# BFS queue: (entity_id, current_depth)
queue: deque[tuple[str, int]] = deque()
visited: set[str] = set()
# Initialize queue with starting entities
for entity_id in entity_ids:
cursor = conn.execute(
"SELECT * FROM entities WHERE id = ?",
(entity_id,),
)
row = cursor.fetchone()
if row is not None:
queue.append((entity_id, 0))
visited.add(entity_id)
collected_entities[entity_id] = self._row_to_entity(row)
# BFS traversal
while queue:
current_id, depth = queue.popleft()
if depth >= max_hops:
continue
# Build relationship query based on direction
if direction == RelationshipDirection.OUTGOING:
rel_query = "SELECT * FROM relationships WHERE source_id = ?"
rel_params: list[Any] = [current_id]
elif direction == RelationshipDirection.INCOMING:
rel_query = "SELECT * FROM relationships WHERE target_id = ?"
rel_params = [current_id]
else: # BOTH
rel_query = (
"SELECT * FROM relationships WHERE source_id = ? OR target_id = ?"
)
rel_params = [current_id, current_id]
# Filter by relation types if specified
if relation_types is not None and len(relation_types) > 0:
placeholders = ", ".join("?" * len(relation_types))
rel_query += f" AND relation_type IN ({placeholders})"
rel_params.extend(relation_types)
cursor = conn.execute(rel_query, rel_params)
for rel_row in cursor:
rel = self._row_to_relationship(rel_row)
# Add relationship
collected_relationships[rel.id] = rel
# Determine neighbor based on direction
neighbor_id = None
if direction == RelationshipDirection.OUTGOING:
if rel.source_id == current_id:
neighbor_id = rel.target_id
elif direction == RelationshipDirection.INCOMING:
if rel.target_id == current_id:
neighbor_id = rel.source_id
else: # BOTH
if rel.source_id == current_id:
neighbor_id = rel.target_id
elif rel.target_id == current_id:
neighbor_id = rel.source_id
if neighbor_id is not None and neighbor_id not in visited:
# Fetch neighbor entity
neighbor_cursor = conn.execute(
"SELECT * FROM entities WHERE id = ?",
(neighbor_id,),
)
neighbor_row = neighbor_cursor.fetchone()
if neighbor_row is not None:
visited.add(neighbor_id)
collected_entities[neighbor_id] = self._row_to_entity(neighbor_row)
queue.append((neighbor_id, depth + 1))
return Subgraph(
entities=list(collected_entities.values()),
relationships=list(collected_relationships.values()),
root_entity_ids=entity_ids,
)
async def find_path(
self,
source_id: str,
target_id: str,
max_depth: int = 10,
direction: RelationshipDirection = RelationshipDirection.BOTH,
) -> list[str] | None:
"""Find the shortest path between two entities using BFS.
Args:
source_id: Starting entity ID.
target_id: Target entity ID.
max_depth: Maximum path length to search (default 10).
direction: Direction of relationship traversal.
Returns:
List of entity IDs representing the path (including source and target),
or None if no path exists within max_depth.
"""
with self._lock:
with self._get_conn() as conn:
# Edge cases
if source_id == target_id:
cursor = conn.execute(
"SELECT id FROM entities WHERE id = ?",
(source_id,),
)
return [source_id] if cursor.fetchone() else None
# Check both exist
cursor = conn.execute(
"SELECT id FROM entities WHERE id IN (?, ?)",
(source_id, target_id),
)
found_ids = {row["id"] for row in cursor}
if source_id not in found_ids or target_id not in found_ids:
return None
# BFS with path tracking
queue: deque[tuple[str, list[str]]] = deque()
visited: set[str] = set()
queue.append((source_id, [source_id]))
visited.add(source_id)
while queue:
current_id, path = queue.popleft()
if len(path) > max_depth:
continue
# Build relationship query
if direction == RelationshipDirection.OUTGOING:
rel_query = "SELECT * FROM relationships WHERE source_id = ?"
rel_params = [current_id]
elif direction == RelationshipDirection.INCOMING:
rel_query = "SELECT * FROM relationships WHERE target_id = ?"
rel_params = [current_id]
else:
rel_query = (
"SELECT * FROM relationships WHERE source_id = ? OR target_id = ?"
)
rel_params = [current_id, current_id]
cursor = conn.execute(rel_query, rel_params)
for rel_row in cursor:
# Determine neighbor
neighbor_id = None
if direction == RelationshipDirection.OUTGOING:
if rel_row["source_id"] == current_id:
neighbor_id = rel_row["target_id"]
elif direction == RelationshipDirection.INCOMING:
if rel_row["target_id"] == current_id:
neighbor_id = rel_row["source_id"]
else:
if rel_row["source_id"] == current_id:
neighbor_id = rel_row["target_id"]
elif rel_row["target_id"] == current_id:
neighbor_id = rel_row["source_id"]
if neighbor_id is None or neighbor_id in visited:
continue
new_path = path + [neighbor_id]
if neighbor_id == target_id:
return new_path
if len(new_path) <= max_depth:
visited.add(neighbor_id)
queue.append((neighbor_id, new_path))
return None
# =========================================================================
# User Management Operations
# =========================================================================
async def clear_user(self, user_id: str) -> tuple[int, int]:
"""Clear all entities and relationships for a user.
Args:
user_id: The user ID to clear data for.
Returns:
Tuple of (entities_deleted, relationships_deleted).
"""
with self._lock:
with self._get_conn() as conn:
# Delete relationships for user first (also cascade-deleted)
cursor = conn.execute(
"DELETE FROM relationships WHERE user_id = ?",
(user_id,),
)
relationships_deleted = cursor.rowcount
# Delete entities (cascades remaining relationships)
cursor = conn.execute(
"DELETE FROM entities WHERE user_id = ?",
(user_id,),
)
entities_deleted = cursor.rowcount
conn.commit()
return entities_deleted, relationships_deleted
# =========================================================================
# Utility Methods
# =========================================================================
async def get_entities_for_user(self, user_id: str) -> list[Entity]:
"""Get all entities for a user.
Args:
user_id: The user ID to get entities for.
Returns:
List of all entities belonging to the user.
"""
with self._lock:
with self._get_conn() as conn:
cursor = conn.execute(
"SELECT * FROM entities WHERE user_id = ?",
(user_id,),
)
return [self._row_to_entity(row) for row in cursor]
async def clear(self) -> None:
"""Clear all data from the store."""
with self._lock:
with self._get_conn() as conn:
conn.execute("DELETE FROM relationships")
conn.execute("DELETE FROM entities")
conn.commit()
@property
def entity_count(self) -> int:
"""Get the total number of entities."""
with self._lock:
with self._get_conn() as conn:
cursor = conn.execute("SELECT COUNT(*) FROM entities")
result = cursor.fetchone()[0]
return int(result)
@property
def relationship_count(self) -> int:
"""Get the total number of relationships."""
with self._lock:
with self._get_conn() as conn:
cursor = conn.execute("SELECT COUNT(*) FROM relationships")
result = cursor.fetchone()[0]
return int(result)
def stats(self) -> dict:
"""Get store statistics.
Returns:
Dict with counts and database info.
"""
with self._lock:
with self._get_conn() as conn:
entity_count = conn.execute("SELECT COUNT(*) FROM entities").fetchone()[0]
rel_count = conn.execute("SELECT COUNT(*) FROM relationships").fetchone()[0]
users_count = conn.execute(
"SELECT COUNT(DISTINCT user_id) FROM entities"
).fetchone()[0]
# Get database file size
db_size = self.db_path.stat().st_size if self.db_path.exists() else 0
return {
"entity_count": entity_count,
"relationship_count": rel_count,
"users_count": users_count,
"db_path": str(self.db_path),
"db_size_bytes": db_size,
"page_cache_size_kb": self._page_cache_size_kb,
}
def get_memory_stats(self) -> ComponentStats:
"""Get memory statistics for the MemoryTracker.
Note: SQLite manages its own memory via page cache. We report
the configured cache size plus estimated Python overhead.
Returns:
ComponentStats with current memory usage.
"""
import sys
from ..tracker import ComponentStats
with self._lock:
with self._get_conn() as conn:
entity_count = conn.execute("SELECT COUNT(*) FROM entities").fetchone()[0]
rel_count = conn.execute("SELECT COUNT(*) FROM relationships").fetchone()[0]
# Estimate Python overhead (connection objects, etc.)
# The actual data is on disk, managed by SQLite's page cache
python_overhead = sys.getsizeof(self) + sys.getsizeof(self._lock)
# Page cache memory (this is the bounded amount)
page_cache_bytes = self._page_cache_size_kb * 1024
return ComponentStats(
name="sqlite_graph_store",
entry_count=entity_count + rel_count,
size_bytes=python_overhead + page_cache_bytes,
budget_bytes=page_cache_bytes, # Cache size is the budget
hits=0,
misses=0,
evictions=0,
)
def vacuum(self) -> None:
"""Reclaim unused space in the database file.
Call this periodically after many deletes to reduce file size.
"""
with self._lock:
with self._get_conn() as conn:
conn.execute("VACUUM")
def close(self) -> None:
"""Close any open connections (cleanup).
Note: This store uses connection-per-request pattern,
so there's typically nothing to close.
"""
pass

View file

@ -4,7 +4,7 @@ Provides a fully local memory backend using embedded databases:
- SQLite for memory storage
- HNSW for vector search
- FTS5 for text search
- In-memory graph for relationships
- SQLite graph for relationships (bounded memory, persistent)
No network calls required, fast startup, suitable for development and
single-process production deployments.
@ -12,37 +12,48 @@ single-process production deployments.
from __future__ import annotations
import logging
import uuid
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any
from headroom.memory.adapters.graph import InMemoryGraphStore
from headroom.memory.adapters.graph_models import Entity, Relationship, Subgraph
from headroom.memory.models import Memory
from headroom.memory.ports import MemorySearchResult
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.
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: Whether to persist graph to SQLite (not yet implemented).
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_model: str = "all-MiniLM-L6-v2"
vector_dimension: int = 384
graph_persist: bool = True # Persist graph to SQLite (future feature)
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
@ -97,13 +108,14 @@ class LocalBackend:
self._config = config or LocalBackendConfig()
self._initialized = False
self._hierarchical_memory: HierarchicalMemory | None = None
self._graph: InMemoryGraphStore | None = None
self._graph: InMemoryGraphStore | SQLiteGraphStore | None = None
async def _ensure_initialized(self) -> None:
"""Ensure the backend is initialized with all components.
Creates the HierarchicalMemory system and InMemoryGraphStore
on first use.
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.
"""
if not self._initialized:
from headroom.memory import HierarchicalMemory, MemoryConfig
@ -117,7 +129,33 @@ class LocalBackend:
)
self._hierarchical_memory = await HierarchicalMemory.create(mem_config)
self._graph = InMemoryGraphStore()
# 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
# =========================================================================
@ -520,7 +558,7 @@ class LocalBackend:
"""Whether this backend supports knowledge graph operations.
Returns:
True, as this backend uses InMemoryGraphStore.
True, as this backend uses SQLiteGraphStore (default) or InMemoryGraphStore.
"""
return True
@ -546,11 +584,12 @@ class LocalBackend:
# Graph Operations
# =========================================================================
async def get_graph(self) -> InMemoryGraphStore:
async def get_graph(self) -> InMemoryGraphStore | SQLiteGraphStore:
"""Get the underlying graph store.
Returns:
The InMemoryGraphStore instance.
The graph store instance (SQLiteGraphStore if graph_persist=True,
InMemoryGraphStore otherwise).
"""
await self._ensure_initialized()
assert self._graph is not None

View file

@ -0,0 +1,903 @@
"""Comprehensive integration tests for SQLiteGraphStore.
Tests verify:
- Entity CRUD operations with persistence
- Relationship CRUD operations with CASCADE delete
- Case-insensitive name lookups
- BFS subgraph traversal
- Shortest path finding
- User data isolation
- Memory bounding via page cache
- Database file persistence across instances
"""
from __future__ import annotations
import os
import tempfile
import pytest
from headroom.memory.adapters.graph_models import (
Entity,
Relationship,
RelationshipDirection,
)
from headroom.memory.adapters.sqlite_graph import SQLiteGraphStore
class TestSQLiteGraphStoreEntityOperations:
"""Tests for entity CRUD operations."""
@pytest.fixture
def store(self):
"""Create a temporary SQLite graph store."""
with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
db_path = f.name
store = SQLiteGraphStore(db_path=db_path)
yield store
# Cleanup
if os.path.exists(db_path):
os.unlink(db_path)
@pytest.mark.asyncio
async def test_add_and_get_entity(self, store):
"""Test adding and retrieving an entity."""
entity = Entity(
user_id="user1",
name="Project Alpha",
entity_type="project",
description="A test project",
properties={"priority": "high"},
)
await store.add_entity(entity)
retrieved = await store.get_entity(entity.id)
assert retrieved is not None
assert retrieved.id == entity.id
assert retrieved.user_id == "user1"
assert retrieved.name == "Project Alpha"
assert retrieved.entity_type == "project"
assert retrieved.description == "A test project"
assert retrieved.properties == {"priority": "high"}
@pytest.mark.asyncio
async def test_get_entity_not_found(self, store):
"""Test retrieving a non-existent entity."""
result = await store.get_entity("nonexistent-id")
assert result is None
@pytest.mark.asyncio
async def test_get_entity_by_name_case_insensitive(self, store):
"""Test case-insensitive entity lookup by name."""
entity = Entity(
user_id="user1",
name="MyEntity",
entity_type="test",
)
await store.add_entity(entity)
# All these should find the same entity
assert (await store.get_entity_by_name("user1", "MyEntity")) is not None
assert (await store.get_entity_by_name("user1", "myentity")) is not None
assert (await store.get_entity_by_name("user1", "MYENTITY")) is not None
assert (await store.get_entity_by_name("user1", "mYeNtItY")) is not None
@pytest.mark.asyncio
async def test_get_entity_by_name_user_isolation(self, store):
"""Test that entity lookup is scoped to user."""
entity1 = Entity(user_id="user1", name="SharedName", entity_type="test")
entity2 = Entity(user_id="user2", name="SharedName", entity_type="test")
await store.add_entity(entity1)
await store.add_entity(entity2)
result1 = await store.get_entity_by_name("user1", "SharedName")
result2 = await store.get_entity_by_name("user2", "SharedName")
assert result1 is not None
assert result2 is not None
assert result1.id != result2.id
assert result1.user_id == "user1"
assert result2.user_id == "user2"
@pytest.mark.asyncio
async def test_update_entity(self, store):
"""Test updating an existing entity."""
entity = Entity(
user_id="user1",
name="Original Name",
entity_type="test",
)
await store.add_entity(entity)
# Update the entity
entity.name = "Updated Name"
entity.entity_type = "updated_type"
entity.properties = {"new_key": "new_value"}
await store.add_entity(entity)
retrieved = await store.get_entity(entity.id)
assert retrieved is not None
assert retrieved.name == "Updated Name"
assert retrieved.entity_type == "updated_type"
assert retrieved.properties == {"new_key": "new_value"}
@pytest.mark.asyncio
async def test_delete_entity(self, store):
"""Test deleting an entity."""
entity = Entity(user_id="user1", name="ToDelete", entity_type="test")
await store.add_entity(entity)
assert await store.get_entity(entity.id) is not None
result = await store.delete_entity(entity.id)
assert result is True
assert await store.get_entity(entity.id) is None
@pytest.mark.asyncio
async def test_delete_entity_not_found(self, store):
"""Test deleting a non-existent entity."""
result = await store.delete_entity("nonexistent-id")
assert result is False
class TestSQLiteGraphStoreRelationshipOperations:
"""Tests for relationship CRUD operations."""
@pytest.fixture
async def store_with_entities(self):
"""Create a store with some entities."""
with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
db_path = f.name
store = SQLiteGraphStore(db_path=db_path)
# Create entities
alice = Entity(user_id="user1", name="Alice", entity_type="person")
bob = Entity(user_id="user1", name="Bob", entity_type="person")
charlie = Entity(user_id="user1", name="Charlie", entity_type="person")
await store.add_entity(alice)
await store.add_entity(bob)
await store.add_entity(charlie)
yield store, alice, bob, charlie
# Cleanup
if os.path.exists(db_path):
os.unlink(db_path)
@pytest.mark.asyncio
async def test_add_and_get_relationship(self, store_with_entities):
"""Test adding and retrieving a relationship."""
store, alice, bob, _ = store_with_entities
rel = Relationship(
user_id="user1",
source_id=alice.id,
target_id=bob.id,
relation_type="knows",
weight=0.9,
properties={"since": "2020"},
)
await store.add_relationship(rel)
# Get outgoing relationships from Alice
rels = await store.get_relationships(alice.id, RelationshipDirection.OUTGOING)
assert len(rels) == 1
assert rels[0].id == rel.id
assert rels[0].source_id == alice.id
assert rels[0].target_id == bob.id
assert rels[0].relation_type == "knows"
assert rels[0].weight == 0.9
assert rels[0].properties == {"since": "2020"}
@pytest.mark.asyncio
async def test_get_relationships_by_direction(self, store_with_entities):
"""Test getting relationships by direction."""
store, alice, bob, charlie = store_with_entities
# Alice -> Bob
rel1 = Relationship(
user_id="user1",
source_id=alice.id,
target_id=bob.id,
relation_type="knows",
)
# Charlie -> Alice
rel2 = Relationship(
user_id="user1",
source_id=charlie.id,
target_id=alice.id,
relation_type="follows",
)
await store.add_relationship(rel1)
await store.add_relationship(rel2)
# Outgoing from Alice
outgoing = await store.get_relationships(alice.id, RelationshipDirection.OUTGOING)
assert len(outgoing) == 1
assert outgoing[0].target_id == bob.id
# Incoming to Alice
incoming = await store.get_relationships(alice.id, RelationshipDirection.INCOMING)
assert len(incoming) == 1
assert incoming[0].source_id == charlie.id
# Both directions
both = await store.get_relationships(alice.id, RelationshipDirection.BOTH)
assert len(both) == 2
@pytest.mark.asyncio
async def test_get_relationships_filter_by_type(self, store_with_entities):
"""Test filtering relationships by type."""
store, alice, bob, charlie = store_with_entities
rel1 = Relationship(
user_id="user1",
source_id=alice.id,
target_id=bob.id,
relation_type="knows",
)
rel2 = Relationship(
user_id="user1",
source_id=alice.id,
target_id=charlie.id,
relation_type="manages",
)
await store.add_relationship(rel1)
await store.add_relationship(rel2)
# Filter by type
knows_rels = await store.get_relationships(
alice.id, RelationshipDirection.OUTGOING, relation_type="knows"
)
assert len(knows_rels) == 1
assert knows_rels[0].target_id == bob.id
manages_rels = await store.get_relationships(
alice.id, RelationshipDirection.OUTGOING, relation_type="manages"
)
assert len(manages_rels) == 1
assert manages_rels[0].target_id == charlie.id
@pytest.mark.asyncio
async def test_delete_relationship(self, store_with_entities):
"""Test deleting a relationship."""
store, alice, bob, _ = store_with_entities
rel = Relationship(
user_id="user1",
source_id=alice.id,
target_id=bob.id,
relation_type="knows",
)
await store.add_relationship(rel)
rels = await store.get_relationships(alice.id, RelationshipDirection.OUTGOING)
assert len(rels) == 1
result = await store.delete_relationship(rel.id)
assert result is True
rels = await store.get_relationships(alice.id, RelationshipDirection.OUTGOING)
assert len(rels) == 0
@pytest.mark.asyncio
async def test_cascade_delete_relationships(self, store_with_entities):
"""Test that deleting an entity cascades to its relationships."""
store, alice, bob, charlie = store_with_entities
# Create relationships
rel1 = Relationship(
user_id="user1",
source_id=alice.id,
target_id=bob.id,
relation_type="knows",
)
rel2 = Relationship(
user_id="user1",
source_id=charlie.id,
target_id=alice.id,
relation_type="follows",
)
await store.add_relationship(rel1)
await store.add_relationship(rel2)
# Delete Alice - should cascade delete both relationships
await store.delete_entity(alice.id)
# Both relationships should be gone
bob_rels = await store.get_relationships(bob.id, RelationshipDirection.BOTH)
assert len(bob_rels) == 0
charlie_rels = await store.get_relationships(charlie.id, RelationshipDirection.BOTH)
assert len(charlie_rels) == 0
class TestSQLiteGraphStoreTraversal:
"""Tests for graph traversal operations."""
@pytest.fixture
async def store_with_graph(self):
"""Create a store with a connected graph.
Graph structure:
A -> B -> D
| |
v v
C -> E
"""
with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
db_path = f.name
store = SQLiteGraphStore(db_path=db_path)
# Create entities
a = Entity(user_id="user1", name="A", entity_type="node")
b = Entity(user_id="user1", name="B", entity_type="node")
c = Entity(user_id="user1", name="C", entity_type="node")
d = Entity(user_id="user1", name="D", entity_type="node")
e = Entity(user_id="user1", name="E", entity_type="node")
for entity in [a, b, c, d, e]:
await store.add_entity(entity)
# Create relationships: A->B, A->C, B->D, B->E, C->E
rels = [
Relationship(user_id="user1", source_id=a.id, target_id=b.id, relation_type="edge"),
Relationship(user_id="user1", source_id=a.id, target_id=c.id, relation_type="edge"),
Relationship(user_id="user1", source_id=b.id, target_id=d.id, relation_type="edge"),
Relationship(user_id="user1", source_id=b.id, target_id=e.id, relation_type="edge"),
Relationship(user_id="user1", source_id=c.id, target_id=e.id, relation_type="edge"),
]
for rel in rels:
await store.add_relationship(rel)
yield store, {"A": a, "B": b, "C": c, "D": d, "E": e}
# Cleanup
if os.path.exists(db_path):
os.unlink(db_path)
@pytest.mark.asyncio
async def test_query_subgraph_single_hop(self, store_with_graph):
"""Test querying subgraph with single hop."""
store, nodes = store_with_graph
subgraph = await store.query_subgraph(
[nodes["A"].id], max_hops=1, direction=RelationshipDirection.OUTGOING
)
entity_names = {e.name for e in subgraph.entities}
assert entity_names == {"A", "B", "C"}
assert len(subgraph.relationships) == 2
@pytest.mark.asyncio
async def test_query_subgraph_two_hops(self, store_with_graph):
"""Test querying subgraph with two hops."""
store, nodes = store_with_graph
subgraph = await store.query_subgraph(
[nodes["A"].id], max_hops=2, direction=RelationshipDirection.OUTGOING
)
entity_names = {e.name for e in subgraph.entities}
assert entity_names == {"A", "B", "C", "D", "E"}
assert len(subgraph.relationships) == 5
@pytest.mark.asyncio
async def test_query_subgraph_incoming(self, store_with_graph):
"""Test querying subgraph with incoming direction."""
store, nodes = store_with_graph
subgraph = await store.query_subgraph(
[nodes["E"].id], max_hops=2, direction=RelationshipDirection.INCOMING
)
entity_names = {e.name for e in subgraph.entities}
# E <- B <- A, E <- C <- A
assert "E" in entity_names
assert "B" in entity_names
assert "C" in entity_names
assert "A" in entity_names
@pytest.mark.asyncio
async def test_find_path_direct(self, store_with_graph):
"""Test finding a direct path."""
store, nodes = store_with_graph
path = await store.find_path(nodes["A"].id, nodes["B"].id)
assert path is not None
assert len(path) == 2
assert path[0] == nodes["A"].id
assert path[1] == nodes["B"].id
@pytest.mark.asyncio
async def test_find_path_multi_hop(self, store_with_graph):
"""Test finding a multi-hop path."""
store, nodes = store_with_graph
path = await store.find_path(nodes["A"].id, nodes["D"].id)
assert path is not None
assert len(path) == 3 # A -> B -> D
assert path[0] == nodes["A"].id
assert path[-1] == nodes["D"].id
@pytest.mark.asyncio
async def test_find_path_not_found(self, store_with_graph):
"""Test that None is returned when no path exists."""
store, nodes = store_with_graph
# D has no outgoing edges, so no path from D to A
path = await store.find_path(
nodes["D"].id, nodes["A"].id, direction=RelationshipDirection.OUTGOING
)
assert path is None
@pytest.mark.asyncio
async def test_find_path_self(self, store_with_graph):
"""Test finding a path to self."""
store, nodes = store_with_graph
path = await store.find_path(nodes["A"].id, nodes["A"].id)
assert path is not None
assert path == [nodes["A"].id]
class TestSQLiteGraphStoreUserManagement:
"""Tests for user data management."""
@pytest.fixture
def store(self):
"""Create a temporary SQLite graph store."""
with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
db_path = f.name
store = SQLiteGraphStore(db_path=db_path)
yield store
if os.path.exists(db_path):
os.unlink(db_path)
@pytest.mark.asyncio
async def test_get_entities_for_user(self, store):
"""Test getting all entities for a specific user."""
# Create entities for two users
for i in range(3):
await store.add_entity(
Entity(user_id="user1", name=f"User1Entity{i}", entity_type="test")
)
for i in range(2):
await store.add_entity(
Entity(user_id="user2", name=f"User2Entity{i}", entity_type="test")
)
user1_entities = await store.get_entities_for_user("user1")
user2_entities = await store.get_entities_for_user("user2")
assert len(user1_entities) == 3
assert len(user2_entities) == 2
assert all(e.user_id == "user1" for e in user1_entities)
assert all(e.user_id == "user2" for e in user2_entities)
@pytest.mark.asyncio
async def test_clear_user(self, store):
"""Test clearing all data for a user."""
# Create entities and relationships for two users
alice = Entity(user_id="user1", name="Alice", entity_type="person")
bob = Entity(user_id="user1", name="Bob", entity_type="person")
carol = Entity(user_id="user2", name="Carol", entity_type="person")
await store.add_entity(alice)
await store.add_entity(bob)
await store.add_entity(carol)
rel = Relationship(
user_id="user1",
source_id=alice.id,
target_id=bob.id,
relation_type="knows",
)
await store.add_relationship(rel)
# Clear user1
entities_deleted, rels_deleted = await store.clear_user("user1")
assert entities_deleted == 2
assert rels_deleted == 1
# User2 data should still exist
assert await store.get_entity(carol.id) is not None
assert store.entity_count == 1
@pytest.mark.asyncio
async def test_clear_all(self, store):
"""Test clearing all data."""
# Add some data
for i in range(5):
await store.add_entity(
Entity(user_id=f"user{i}", name=f"Entity{i}", entity_type="test")
)
assert store.entity_count == 5
await store.clear()
assert store.entity_count == 0
assert store.relationship_count == 0
class TestSQLiteGraphStorePersistence:
"""Tests for data persistence across store instances."""
@pytest.mark.asyncio
async def test_data_persists_across_instances(self):
"""Test that data survives store restart."""
with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
db_path = f.name
try:
# Create store and add data
store1 = SQLiteGraphStore(db_path=db_path)
entity = Entity(user_id="user1", name="Persistent", entity_type="test")
await store1.add_entity(entity)
entity_id = entity.id
# Create new store instance pointing to same database
store2 = SQLiteGraphStore(db_path=db_path)
# Data should be there
retrieved = await store2.get_entity(entity_id)
assert retrieved is not None
assert retrieved.name == "Persistent"
finally:
if os.path.exists(db_path):
os.unlink(db_path)
@pytest.mark.asyncio
async def test_relationships_persist(self):
"""Test that relationships persist across instances."""
with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
db_path = f.name
try:
store1 = SQLiteGraphStore(db_path=db_path)
alice = Entity(user_id="user1", name="Alice", entity_type="person")
bob = Entity(user_id="user1", name="Bob", entity_type="person")
await store1.add_entity(alice)
await store1.add_entity(bob)
rel = Relationship(
user_id="user1",
source_id=alice.id,
target_id=bob.id,
relation_type="knows",
)
await store1.add_relationship(rel)
# New instance
store2 = SQLiteGraphStore(db_path=db_path)
rels = await store2.get_relationships(alice.id, RelationshipDirection.OUTGOING)
assert len(rels) == 1
assert rels[0].target_id == bob.id
finally:
if os.path.exists(db_path):
os.unlink(db_path)
class TestSQLiteGraphStoreMemoryStats:
"""Tests for memory statistics and bounding."""
@pytest.fixture
def store(self):
"""Create a temporary SQLite graph store."""
with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
db_path = f.name
store = SQLiteGraphStore(db_path=db_path, page_cache_size_kb=4096)
yield store
if os.path.exists(db_path):
os.unlink(db_path)
@pytest.mark.asyncio
async def test_stats(self, store):
"""Test getting store statistics."""
# Add some data
for i in range(10):
await store.add_entity(Entity(user_id="user1", name=f"Entity{i}", entity_type="test"))
stats = store.stats()
assert stats["entity_count"] == 10
assert stats["relationship_count"] == 0
assert stats["users_count"] == 1
assert stats["page_cache_size_kb"] == 4096
assert "db_path" in stats
assert stats["db_size_bytes"] > 0
@pytest.mark.asyncio
async def test_memory_stats(self, store):
"""Test memory statistics for MemoryTracker."""
# Add some data
for i in range(5):
await store.add_entity(Entity(user_id="user1", name=f"Entity{i}", entity_type="test"))
stats = store.get_memory_stats()
assert stats.name == "sqlite_graph_store"
assert stats.entry_count == 5
assert stats.size_bytes > 0
# Budget should be the page cache size
assert stats.budget_bytes == 4096 * 1024
@pytest.mark.asyncio
async def test_vacuum(self, store):
"""Test vacuuming the database."""
# Add and delete data
entities = []
for i in range(100):
e = Entity(user_id="user1", name=f"Entity{i}", entity_type="test")
await store.add_entity(e)
entities.append(e)
# Delete all
for e in entities:
await store.delete_entity(e.id)
# Get size before vacuum
stats_before = store.stats()
# Vacuum
store.vacuum()
# Size should be smaller or same after vacuum
stats_after = store.stats()
assert stats_after["db_size_bytes"] <= stats_before["db_size_bytes"]
class TestSQLiteGraphStoreEdgeCases:
"""Tests for edge cases and error handling."""
@pytest.fixture
def store(self):
"""Create a temporary SQLite graph store."""
with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
db_path = f.name
store = SQLiteGraphStore(db_path=db_path)
yield store
if os.path.exists(db_path):
os.unlink(db_path)
@pytest.mark.asyncio
async def test_empty_subgraph_query(self, store):
"""Test querying subgraph with no entities."""
subgraph = await store.query_subgraph(["nonexistent-id"])
assert len(subgraph.entities) == 0
assert len(subgraph.relationships) == 0
@pytest.mark.asyncio
async def test_entity_with_special_characters(self, store):
"""Test entity names with special characters."""
entity = Entity(
user_id="user1",
name='Test\'s "Entity" (Special & <More>)',
entity_type="test",
description="Description with 'quotes' and \"more\"",
properties={"key": "value with 'quotes'"},
)
await store.add_entity(entity)
retrieved = await store.get_entity(entity.id)
assert retrieved is not None
assert retrieved.name == entity.name
assert retrieved.description == entity.description
assert retrieved.properties == entity.properties
@pytest.mark.asyncio
async def test_entity_with_unicode(self, store):
"""Test entity names with unicode characters."""
entity = Entity(
user_id="user1",
name="Test 你好 🚀 Ñoño",
entity_type="test",
)
await store.add_entity(entity)
retrieved = await store.get_entity(entity.id)
assert retrieved is not None
assert retrieved.name == "Test 你好 🚀 Ñoño"
@pytest.mark.asyncio
async def test_large_properties(self, store):
"""Test entities with large properties."""
large_props = {f"key_{i}": f"value_{i}" * 100 for i in range(100)}
entity = Entity(
user_id="user1",
name="LargeEntity",
entity_type="test",
properties=large_props,
)
await store.add_entity(entity)
retrieved = await store.get_entity(entity.id)
assert retrieved is not None
assert retrieved.properties == large_props
@pytest.mark.asyncio
async def test_concurrent_access(self, store):
"""Test basic concurrent access (thread-safe pattern)."""
import asyncio
async def add_entity(i: int):
e = Entity(user_id="user1", name=f"Concurrent{i}", entity_type="test")
await store.add_entity(e)
return e.id
# Add entities concurrently
ids = await asyncio.gather(*[add_entity(i) for i in range(20)])
# All should exist
assert len(ids) == 20
assert store.entity_count == 20
class TestSQLiteGraphStoreMemoryTrackerIntegration:
"""Tests for MemoryTracker integration."""
@pytest.fixture
def store(self):
"""Create a temporary SQLite graph store."""
with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
db_path = f.name
store = SQLiteGraphStore(db_path=db_path, page_cache_size_kb=4096)
yield store
if os.path.exists(db_path):
os.unlink(db_path)
@pytest.mark.asyncio
async def test_memory_tracker_registration(self, store):
"""Test registering SQLiteGraphStore with MemoryTracker."""
from headroom.memory.tracker import MemoryTracker
tracker = MemoryTracker.get()
# Unregister if already registered from previous test
tracker.unregister("sqlite_graph_test")
# Register the store
tracker.register("sqlite_graph_test", store.get_memory_stats)
try:
# Add some data
for i in range(10):
await store.add_entity(
Entity(user_id="user1", name=f"Entity{i}", entity_type="test")
)
# Get report
report = tracker.get_report()
# Find our component in the report's components dict
assert "sqlite_graph_test" in report.components
graph_stats = report.components["sqlite_graph_test"]
assert graph_stats is not None
assert graph_stats.entry_count == 10
assert graph_stats.budget_bytes == 4096 * 1024 # 4MB cache
finally:
tracker.unregister("sqlite_graph_test")
@pytest.mark.asyncio
async def test_memory_stats_tracks_growth(self, store):
"""Test that memory stats track entity/relationship growth."""
stats_before = store.get_memory_stats()
assert stats_before.entry_count == 0
# Add entities
entities = []
for i in range(5):
e = Entity(user_id="user1", name=f"Entity{i}", entity_type="test")
await store.add_entity(e)
entities.append(e)
stats_after_entities = store.get_memory_stats()
assert stats_after_entities.entry_count == 5
# Add relationships
for i in range(4):
rel = Relationship(
user_id="user1",
source_id=entities[i].id,
target_id=entities[i + 1].id,
relation_type="connected",
)
await store.add_relationship(rel)
stats_after_rels = store.get_memory_stats()
assert stats_after_rels.entry_count == 9 # 5 entities + 4 relationships
@pytest.mark.asyncio
async def test_memory_stats_bounded_by_cache(self, store):
"""Test that reported size is bounded by cache size."""
# Add many entities
for i in range(100):
await store.add_entity(Entity(user_id="user1", name=f"Entity{i}", entity_type="test"))
stats = store.get_memory_stats()
# Size should be bounded by cache + overhead
# Cache is 4MB, overhead should be small
assert stats.size_bytes < 5 * 1024 * 1024 # Less than 5MB
assert stats.budget_bytes == 4096 * 1024 # Exactly 4MB
@pytest.mark.asyncio
async def test_memory_report_includes_sqlite_graph(self, store):
"""Test that MemoryTracker report includes SQLiteGraphStore stats."""
from headroom.memory.tracker import MemoryTracker
tracker = MemoryTracker.get()
# Unregister if already registered from previous test
tracker.unregister("graph_report_test")
tracker.register("graph_report_test", store.get_memory_stats)
try:
# Add data
await store.add_entity(Entity(user_id="user1", name="Test", entity_type="test"))
# Get full report dict
report = tracker.get_report()
report_dict = report.to_dict()
# Verify structure
assert "components" in report_dict
assert len(report_dict["components"]) > 0
# Find our component (components is a dict in to_dict output)
assert "graph_report_test" in report_dict["components"]
comp = report_dict["components"]["graph_report_test"]
assert comp["name"] == "sqlite_graph_store"
assert comp["entry_count"] == 1
assert "size_bytes" in comp
assert "budget_bytes" in comp
finally:
tracker.unregister("graph_report_test")
@pytest.mark.asyncio
async def test_memory_stats_after_delete(self, store):
"""Test that memory stats decrease after deletion."""
# Add entities
entities = []
for i in range(10):
e = Entity(user_id="user1", name=f"Entity{i}", entity_type="test")
await store.add_entity(e)
entities.append(e)
stats_before_delete = store.get_memory_stats()
assert stats_before_delete.entry_count == 10
# Delete half
for e in entities[:5]:
await store.delete_entity(e.id)
stats_after_delete = store.get_memory_stats()
assert stats_after_delete.entry_count == 5