speed up batch writes

hindsight-api/hindsight_api/engine/retain/link_utils.py

@@ -529,25 +529,53 @@ async def create_semantic_links_batch(
         raise
 
 
-async def insert_entity_links_batch(conn, links: List[tuple]):
+async def insert_entity_links_batch(conn, links: List[tuple], chunk_size: int = 5000):
     """
-    Insert all entity links in a single batch.
+    Insert all entity links in bulk using unnest for efficiency.
+
+    Uses PostgreSQL unnest() to insert many rows in a single query,
+    which is much faster than executemany over high-latency connections.
 
     Args:
         conn: Database connection
         links: List of tuples (from_unit_id, to_unit_id, link_type, weight, entity_id)
+        chunk_size: Number of rows per batch (default 5000)
     """
     if not links:
         return
 
-    await conn.executemany(
+    import uuid as uuid_mod
-        """
+
-        INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
+    # Process in chunks to avoid query size limits
-        VALUES ($1, $2, $3, $4, $5)
+    for i in range(0, len(links), chunk_size):
-        ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
+        chunk = links[i:i + chunk_size]
-        """,
+
-        links
+        # Separate into arrays for unnest
-    )
+        from_ids = []
+        to_ids = []
+        link_types = []
+        weights = []
+        entity_ids = []
+
+        for from_id, to_id, link_type, weight, entity_id in chunk:
+            from_ids.append(uuid_mod.UUID(from_id) if isinstance(from_id, str) else from_id)
+            to_ids.append(uuid_mod.UUID(to_id) if isinstance(to_id, str) else to_id)
+            link_types.append(link_type)
+            weights.append(weight)
+            entity_ids.append(
+                uuid_mod.UUID(str(entity_id)) if entity_id and not isinstance(entity_id, uuid_mod.UUID)
+                else entity_id
+            )
+
+        # Use unnest to insert all rows in one query
+        await conn.execute(
+            """
+            INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
+            SELECT * FROM unnest($1::uuid[], $2::uuid[], $3::text[], $4::float[], $5::uuid[])
+            ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
+            """,
+            from_ids, to_ids, link_types, weights, entity_ids
+        )
 
 
 async def create_causal_links_batch(

uv.lock

@@ -1141,7 +1141,7 @@ provides-extras = ["test"]
 
 [[package]]
 name = "hindsight-api"
-version = "0.0.17"
+version = "0.0.18"
 source = { editable = "hindsight-api" }
 dependencies = [
     { name = "alembic" },
@@ -1243,7 +1243,7 @@ dev = [
 
 [[package]]
 name = "hindsight-client"
-version = "0.0.17"
+version = "0.0.18"
 source = { editable = "hindsight-clients/python" }
 dependencies = [
     { name = "aiohttp" },
@@ -1275,7 +1275,7 @@ provides-extras = ["test"]
 
 [[package]]
 name = "hindsight-dev"
-version = "0.0.17"
+version = "0.0.18"
 source = { editable = "hindsight-dev" }
 dependencies = [
     { name = "hindsight-api" },