diff --git a/headroom/memory/adapters/__init__.py b/headroom/memory/adapters/__init__.py index d0dc11bb8..f23eabc9b 100644 --- a/headroom/memory/adapters/__init__.py +++ b/headroom/memory/adapters/__init__.py @@ -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", diff --git a/headroom/memory/adapters/sqlite_graph.py b/headroom/memory/adapters/sqlite_graph.py new file mode 100644 index 000000000..8f14c4ad4 --- /dev/null +++ b/headroom/memory/adapters/sqlite_graph.py @@ -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 diff --git a/headroom/memory/backends/local.py b/headroom/memory/backends/local.py index 47c72df76..cb6eb1235 100644 --- a/headroom/memory/backends/local.py +++ b/headroom/memory/backends/local.py @@ -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 diff --git a/tests/test_sqlite_graph_store.py b/tests/test_sqlite_graph_store.py new file mode 100644 index 000000000..9fb5a90b0 --- /dev/null +++ b/tests/test_sqlite_graph_store.py @@ -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 & )', + 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