graphmemory 1.2.1__tar.gz → 1.3.0__tar.gz

{graphmemory-1.2.1 → graphmemory-1.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: graphmemory
-Version: 1.2.1
+Version: 1.3.0
 Summary: Graph-based memory system using DuckDB
 Project-URL: Homepage, https://github.com/bradAGI/GraphMemory
 Project-URL: Repository, https://github.com/bradAGI/GraphMemory
@@ -224,7 +224,7 @@ All IDs are auto-generated UUIDs. All models are [Pydantic](https://docs.pydanti
 
 | Method | Description |
 |--------|-------------|
-| `GraphMemory(database=None, vector_length=3, distance_metric='l2')` | Initialize. `None` = in-memory. |
+| `GraphMemory(database=None, vector_length=3, distance_metric='l2', hnsw_ef_construction=128, hnsw_ef_search=64, hnsw_m=16, auto_index=True)` | Initialize. `None` = in-memory. HNSW index auto-created. |
 | `close()` | Close connection (thread-safe, idempotent). |
 | `transaction()` | Context manager for atomic operations. |
 
@@ -262,7 +262,8 @@ All IDs are auto-generated UUIDs. All models are [Pydantic](https://docs.pydanti
 | `nearest_nodes(vector, limit) -> list[NearestNode]` | Vector similarity search. |
 | `search_nodes(query_text, limit=10) -> list[SearchResult]` | Full-text BM25 search. |
 | `hybrid_search(query_text, query_vector, ...) -> list[SearchResult]` | Combined text + vector search. |
-| `create_index()` | Create HNSW index for faster vector search. |
+| `create_index(ef_construction=None, ef_search=None, m=None)` | Create/recreate HNSW index with tunable params. Auto-called on init. |
+| `compact_index()` | Compact HNSW index to reclaim space after deletions. |
 
 ### Retrieval
 
@@ -295,7 +296,7 @@ See `examples/` for complete usage:
 
 ## Testing
 
-265 tests covering all functionality.
+291 tests covering all functionality.
 
 ```sh
 python3 -m pytest tests/tests.py -v

{graphmemory-1.2.1 → graphmemory-1.3.0}/README.md

@@ -197,7 +197,7 @@ All IDs are auto-generated UUIDs. All models are [Pydantic](https://docs.pydanti
 
 | Method | Description |
 |--------|-------------|
-| `GraphMemory(database=None, vector_length=3, distance_metric='l2')` | Initialize. `None` = in-memory. |
+| `GraphMemory(database=None, vector_length=3, distance_metric='l2', hnsw_ef_construction=128, hnsw_ef_search=64, hnsw_m=16, auto_index=True)` | Initialize. `None` = in-memory. HNSW index auto-created. |
 | `close()` | Close connection (thread-safe, idempotent). |
 | `transaction()` | Context manager for atomic operations. |
 
@@ -235,7 +235,8 @@ All IDs are auto-generated UUIDs. All models are [Pydantic](https://docs.pydanti
 | `nearest_nodes(vector, limit) -> list[NearestNode]` | Vector similarity search. |
 | `search_nodes(query_text, limit=10) -> list[SearchResult]` | Full-text BM25 search. |
 | `hybrid_search(query_text, query_vector, ...) -> list[SearchResult]` | Combined text + vector search. |
-| `create_index()` | Create HNSW index for faster vector search. |
+| `create_index(ef_construction=None, ef_search=None, m=None)` | Create/recreate HNSW index with tunable params. Auto-called on init. |
+| `compact_index()` | Compact HNSW index to reclaim space after deletions. |
 
 ### Retrieval
 
@@ -268,7 +269,7 @@ See `examples/` for complete usage:
 
 ## Testing
 
-265 tests covering all functionality.
+291 tests covering all functionality.
 
 ```sh
 python3 -m pytest tests/tests.py -v

{graphmemory-1.2.1 → graphmemory-1.3.0}/graphmemory/database.py

@@ -86,7 +86,8 @@ class GraphMemory:
         'inner_product': {'function': 'array_negative_inner_product', 'hnsw_metric': 'ip'},
     }
 
-    def __init__(self, database=None, vector_length=3, distance_metric='l2', max_retries=3, retry_base_delay=0.1):
+    def __init__(self, database=None, vector_length=3, distance_metric='l2', max_retries=3, retry_base_delay=0.1,
+                 hnsw_ef_construction=128, hnsw_ef_search=64, hnsw_m=16, auto_index=True):
         if distance_metric not in self.DISTANCE_METRICS:
             raise ValueError(
                 f"Invalid distance_metric '{distance_metric}'. "
@@ -97,9 +98,13 @@ class GraphMemory:
         self.distance_metric = distance_metric
         self.max_retries = max_retries
         self.retry_base_delay = retry_base_delay
+        self.hnsw_ef_construction = hnsw_ef_construction
+        self.hnsw_ef_search = hnsw_ef_search
+        self.hnsw_m = hnsw_m
         self._lock = threading.RLock()
         self._fts_initialized = False
         self._fts_dirty = True
+        self._hnsw_indexed = False
         self._closed = False
         self.conn = duckdb.connect(database=self.database)
         self._load_vss_extension()
@@ -116,6 +121,9 @@ class GraphMemory:
             self._create_tables()
             logger.info("Tables created or verified successfully.")
 
+        if auto_index:
+            self._ensure_hnsw_index()
+
     def cursor(self):
         """Return a new DuckDB cursor for individual operations.
 
@@ -150,6 +158,8 @@ class GraphMemory:
             self._configure_database()
             self._fts_initialized = False
             self._fts_dirty = True
+            self._hnsw_indexed = False
+            self._ensure_hnsw_index()
             logger.info("Reconnection successful.")
 
     def close(self):
@@ -191,6 +201,8 @@ class GraphMemory:
 
     def set_vector_length(self, vector_length):
         self.vector_length = vector_length
+        self._hnsw_indexed = False
+        self._ensure_hnsw_index()
         logger.info(f"Vector length set to: {self.vector_length}")
 
     def _create_tables(self):
@@ -303,6 +315,7 @@ class GraphMemory:
                 cur.execute(
                     "DELETE FROM nodes WHERE id = ?;", (str(node_id),))
                 self._fts_dirty = True
+            self.compact_index()
         except duckdb.Error as e:
             logger.error(f"Error deleting node: {e}")
 
@@ -321,6 +334,7 @@ class GraphMemory:
                 cur.execute(
                     f"DELETE FROM nodes WHERE id IN ({placeholders});", id_strs)
                 self._fts_dirty = True
+            self.compact_index()
         except duckdb.Error as e:
             logger.error(f"Error during bulk delete nodes: {e}")
 
@@ -920,15 +934,53 @@ class GraphMemory:
             logger.error(f"Error updating edge: {e}")
             return False
 
+    def _ensure_hnsw_index(self):
+        """Create HNSW index if not already present. Called automatically on init."""
+        if self._hnsw_indexed:
+            return
+        try:
+            nodes_exist = self.conn.execute(
+                "SELECT 1 FROM information_schema.tables WHERE table_name = 'nodes';"
+            ).fetchone()
+            if nodes_exist:
+                self.create_index()
+        except duckdb.Error:
+            pass
+
     @with_retry()
-    def create_index(self):
+    def create_index(self, ef_construction: int | None = None, ef_search: int | None = None, m: int | None = None):
+        """Create or recreate the HNSW vector index.
+
+        Args:
+            ef_construction: Candidate vertices during build (default from init).
+            ef_search: Candidate vertices during search (default from init).
+            m: Max neighbors per vertex (default from init).
+        """
+        ef_c = ef_construction or self.hnsw_ef_construction
+        ef_s = ef_search or self.hnsw_ef_search
+        m_val = m or self.hnsw_m
+        hnsw_metric = self.DISTANCE_METRICS[self.distance_metric]['hnsw_metric']
         with self._lock:
             try:
-                hnsw_metric = self.DISTANCE_METRICS[self.distance_metric]['hnsw_metric']
+                # Drop existing index first to allow metric/param changes
+                self.conn.execute("DROP INDEX IF EXISTS vss_idx;")
                 self.conn.execute(
-                    f"CREATE INDEX IF NOT EXISTS vss_idx ON nodes USING HNSW(vector) WITH (metric = '{hnsw_metric}');")
+                    f"CREATE INDEX vss_idx ON nodes USING HNSW(vector) "
+                    f"WITH (metric = '{hnsw_metric}', ef_construction = {ef_c}, ef_search = {ef_s}, M = {m_val});"
+                )
+                self._hnsw_indexed = True
+                logger.info(f"HNSW index created (metric={hnsw_metric}, ef_construction={ef_c}, ef_search={ef_s}, M={m_val}).")
+            except duckdb.Error as e:
+                logger.error(f"Error creating HNSW index: {e}")
+
+    def compact_index(self):
+        """Compact the HNSW index to reclaim space after deletions."""
+        with self._lock:
+            try:
+                self.conn.execute("PRAGMA hnsw_compact_index('vss_idx');")
+                logger.info("HNSW index compacted.")
             except duckdb.Error as e:
-                logger.error(f"Error creating index: {e}")
+                logger.error(f"Error compacting HNSW index: {e}")
 
     @with_retry()
     def nearest_nodes(self, vector: list[float], limit: int) -> list[NearestNode]:
@@ -1334,9 +1386,10 @@ class GraphMemory:
 
             # Collect vector similarity results
             vss_results = {}
+            dist_func = self.DISTANCE_METRICS[self.distance_metric]['function']
             vss_query = f"""
             SELECT id, type, properties, vector,
-                   array_distance(vector, CAST(? AS FLOAT[{self.vector_length}])) AS distance
+                   {dist_func}(vector, CAST(? AS FLOAT[{self.vector_length}])) AS distance
             FROM nodes
             WHERE vector IS NOT NULL
             ORDER BY distance;

{graphmemory-1.2.1 → graphmemory-1.3.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "graphmemory"
-version = "1.2.1"
+version = "1.3.0"
 description = "Graph-based memory system using DuckDB"
 readme = "README.md"
 license = "MIT"

{graphmemory-1.2.1 → graphmemory-1.3.0}/tests/tests.py

@@ -2729,5 +2729,147 @@ class TestFuzzyMatching(unittest.TestCase):
         self.assertEqual(remaining, 2)
 
 
+class TestHNSWIndex(unittest.TestCase):
+
+    def test_auto_index_on_init(self):
+        db = GraphMemory(database=':memory:', vector_length=3)
+        self.assertTrue(db._hnsw_indexed)
+        db.close()
+
+    def test_auto_index_disabled(self):
+        db = GraphMemory(database=':memory:', vector_length=3, auto_index=False)
+        self.assertFalse(db._hnsw_indexed)
+        db.close()
+
+    def test_create_index_with_custom_params(self):
+        db = GraphMemory(database=':memory:', vector_length=3, auto_index=False)
+        db.create_index(ef_construction=64, ef_search=32, m=8)
+        self.assertTrue(db._hnsw_indexed)
+        db.close()
+
+    def test_create_index_uses_configured_metric(self):
+        for metric in ['l2', 'cosine', 'inner_product']:
+            db = GraphMemory(database=':memory:', vector_length=3, distance_metric=metric)
+            self.assertTrue(db._hnsw_indexed)
+            db.close()
+
+    def test_create_index_idempotent_recreate(self):
+        db = GraphMemory(database=':memory:', vector_length=3)
+        db.create_index()
+        db.create_index()
+        self.assertTrue(db._hnsw_indexed)
+        db.close()
+
+    def test_set_vector_length_rebuilds_index(self):
+        db = GraphMemory(database=':memory:', vector_length=3)
+        self.assertTrue(db._hnsw_indexed)
+        db._hnsw_indexed = False
+        db.set_vector_length(5)
+        self.assertTrue(db._hnsw_indexed)
+        self.assertEqual(db.vector_length, 5)
+        db.close()
+
+    def test_compact_index_no_error(self):
+        db = GraphMemory(database=':memory:', vector_length=3)
+        node = Node(type="Test", properties={"name": "A"}, vector=[1.0, 0.0, 0.0])
+        db.insert_node(node)
+        db.delete_node(node.id)
+        db.compact_index()
+        db.close()
+
+    def test_reconnect_rebuilds_index(self):
+        import tempfile
+        path = tempfile.mktemp(suffix='.db')
+        try:
+            db = GraphMemory(database=path, vector_length=3)
+            db._hnsw_indexed = False
+            db._reconnect()
+            self.assertTrue(db._hnsw_indexed)
+            db.close()
+        finally:
+            if os.path.exists(path):
+                os.unlink(path)
+
+    def test_hnsw_params_stored(self):
+        db = GraphMemory(database=':memory:', vector_length=3,
+                         hnsw_ef_construction=256, hnsw_ef_search=128, hnsw_m=32)
+        self.assertEqual(db.hnsw_ef_construction, 256)
+        self.assertEqual(db.hnsw_ef_search, 128)
+        self.assertEqual(db.hnsw_m, 32)
+        db.close()
+
+
+class TestHybridSearchMetric(unittest.TestCase):
+
+    def setUp(self):
+        self.db = GraphMemory(database=':memory:', vector_length=3, distance_metric='cosine')
+        self.db.insert_node(Node(type="Doc", properties={"text": "machine learning"}, vector=[1.0, 0.0, 0.0]))
+        self.db.insert_node(Node(type="Doc", properties={"text": "deep learning"}, vector=[0.9, 0.1, 0.0]))
+        self.db.insert_node(Node(type="Doc", properties={"text": "cooking recipes"}, vector=[0.0, 0.0, 1.0]))
+
+    def tearDown(self):
+        self.db.close()
+
+    def test_hybrid_search_uses_cosine_metric(self):
+        results = self.db.hybrid_search(
+            query_text="learning",
+            query_vector=[1.0, 0.0, 0.0],
+            limit=3
+        )
+        self.assertGreater(len(results), 0)
+        # The learning docs should score higher than cooking
+        names = [r.node.properties.get("text") for r in results]
+        self.assertIn("machine learning", names[:2])
+
+    def test_hybrid_search_inner_product(self):
+        db = GraphMemory(database=':memory:', vector_length=3, distance_metric='inner_product')
+        db.insert_node(Node(type="Doc", properties={"text": "similar"}, vector=[1.0, 0.0, 0.0]))
+        db.insert_node(Node(type="Doc", properties={"text": "different"}, vector=[0.0, 0.0, 1.0]))
+        results = db.hybrid_search(
+            query_text="similar",
+            query_vector=[1.0, 0.0, 0.0],
+            limit=2
+        )
+        self.assertGreater(len(results), 0)
+        db.close()
+
+    def test_hybrid_search_l2_metric(self):
+        db = GraphMemory(database=':memory:', vector_length=3, distance_metric='l2')
+        db.insert_node(Node(type="Doc", properties={"text": "near"}, vector=[0.1, 0.0, 0.0]))
+        db.insert_node(Node(type="Doc", properties={"text": "far"}, vector=[9.0, 9.0, 9.0]))
+        results = db.hybrid_search(
+            query_text="near",
+            query_vector=[0.0, 0.0, 0.0],
+            limit=2
+        )
+        self.assertGreater(len(results), 0)
+        self.assertEqual(results[0].node.properties["text"], "near")
+        db.close()
+
+
+class TestCompactAfterDelete(unittest.TestCase):
+
+    def setUp(self):
+        self.db = GraphMemory(database=':memory:', vector_length=3)
+
+    def tearDown(self):
+        self.db.close()
+
+    def test_delete_node_compacts(self):
+        node = Node(type="Test", properties={"name": "A"}, vector=[1.0, 0.0, 0.0])
+        self.db.insert_node(node)
+        # Should not raise — compact_index called internally
+        self.db.delete_node(node.id)
+        self.assertEqual(len(self.db.nodes_to_json()), 0)
+
+    def test_bulk_delete_nodes_compacts(self):
+        n1 = Node(type="Test", properties={"name": "A"}, vector=[1.0, 0.0, 0.0])
+        n2 = Node(type="Test", properties={"name": "B"}, vector=[0.0, 1.0, 0.0])
+        self.db.insert_node(n1)
+        self.db.insert_node(n2)
+        self.db.bulk_delete_nodes([n1.id, n2.id])
+        self.assertEqual(len(self.db.nodes_to_json()), 0)
+
+
 if __name__ == '__main__':
     unittest.main()