mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
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:
parent
e16691dd38
commit
f2ecf1d2d0
4 changed files with 1715 additions and 12 deletions
|
|
@ -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",
|
||||
|
|
|
|||
758
headroom/memory/adapters/sqlite_graph.py
Normal file
758
headroom/memory/adapters/sqlite_graph.py
Normal 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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
903
tests/test_sqlite_graph_store.py
Normal file
903
tests/test_sqlite_graph_store.py
Normal 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
|
||||
Loading…
Add table
Add a link
Reference in a new issue