holetron-lab/fleet-memory
1
0
Fork 0
Code Releases Packages Activity Actions

speed up batch writes

Browse source
This commit is contained in:
Nicolò Boschi 2025-12-04 16:44:48 +01:00
parent bb434f3f1a
commit 3402bf15ee
2 changed files with 41 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

48
hindsight-api/hindsight_api/engine/retain/link_utils.py
View file

@ -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(

6
uv.lock
View file

@ -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" },