diff --git a/hindsight-api/hindsight_api/engine/retain/link_utils.py b/hindsight-api/hindsight_api/engine/retain/link_utils.py index e6f21c39..45b60c3d 100644 --- a/hindsight-api/hindsight_api/engine/retain/link_utils.py +++ b/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( - """ - INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id) - VALUES ($1, $2, $3, $4, $5) - ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING - """, - links - ) + import uuid as uuid_mod + + # Process in chunks to avoid query size limits + for i in range(0, len(links), chunk_size): + chunk = links[i:i + chunk_size] + + # 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( diff --git a/uv.lock b/uv.lock index b0eece75..2bebb93f 100644 --- a/uv.lock +++ b/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" },