From 91bc3b02bc9a9f900534b533f8dba88b9a7e9b5d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Thu, 4 Dec 2025 16:52:43 +0100 Subject: [PATCH] speed up batch writes --- hindsight-api/hindsight_api/api/http.py | 12 +-- hindsight-api/hindsight_api/api/mcp.py | 2 +- .../hindsight_api/engine/memory_engine.py | 8 +- .../hindsight_api/engine/response_models.py | 2 +- .../engine/retain/fact_extraction.py | 14 ++-- .../hindsight_api/engine/retain/link_utils.py | 74 +++++++++++-------- .../hindsight_api/engine/retain/types.py | 2 +- hindsight-api/hindsight_api/models.py | 2 +- uv.lock | 6 +- 9 files changed, 66 insertions(+), 56 deletions(-) diff --git a/hindsight-api/hindsight_api/api/http.py b/hindsight-api/hindsight_api/api/http.py index 8ba35431..3f4b6ffa 100644 --- a/hindsight-api/hindsight_api/api/http.py +++ b/hindsight-api/hindsight_api/api/http.py @@ -84,7 +84,7 @@ class RecallRequest(BaseModel): model_config = ConfigDict(json_schema_extra={ "example": { "query": "What did Alice say about machine learning?", - "types": ["world", "bank"], + "types": ["world", "interactions"], "budget": "mid", "max_tokens": 4096, "trace": True, @@ -417,7 +417,7 @@ class ReflectResponse(BaseModel): { "id": "456", "text": "I discussed AI applications last week", - "type": "bank" + "type": "interactions" } ] } @@ -901,7 +901,7 @@ def _register_routes(app: FastAPI): The type parameter is optional and must be one of: - 'world': General knowledge about people, places, events, and things that happen - - 'bank': Memories about what the AI agent did, actions taken, and tasks performed + - 'interactions': Memories about interactions, conversations, actions taken, and tasks performed - 'opinion': The bank's formed beliefs, perspectives, and viewpoints Set include_entities=true to get entity observations alongside recall results. @@ -914,10 +914,10 @@ def _register_routes(app: FastAPI): try: # Validate types - valid_fact_types = ["world", "bank", "opinion"] + valid_fact_types = ["world", "interactions", "opinion"] - # Default to world, agent, opinion if not specified (exclude observation by default) - fact_types = request.types if request.types else ["world", "bank", "opinion"] + # Default to world, interactions, opinion if not specified (exclude observation by default) + fact_types = request.types if request.types else ["world", "interactions", "opinion"] for ft in fact_types: if ft not in valid_fact_types: raise HTTPException( diff --git a/hindsight-api/hindsight_api/api/mcp.py b/hindsight-api/hindsight_api/api/mcp.py index 5ad2db58..6940b427 100644 --- a/hindsight-api/hindsight_api/api/mcp.py +++ b/hindsight-api/hindsight_api/api/mcp.py @@ -90,7 +90,7 @@ def create_mcp_server(memory: MemoryEngine) -> FastMCP: search_result = await memory.recall_async( bank_id=bank_id, query=query, - fact_type=["world", "bank", "opinion"], + fact_type=["world", "interactions", "opinion"], budget=Budget.LOW ) diff --git a/hindsight-api/hindsight_api/engine/memory_engine.py b/hindsight-api/hindsight_api/engine/memory_engine.py index 0a24d440..6b1a183d 100644 --- a/hindsight-api/hindsight_api/engine/memory_engine.py +++ b/hindsight-api/hindsight_api/engine/memory_engine.py @@ -676,7 +676,7 @@ class MemoryEngine: context: Context about when/why this memory was formed event_date: When the event occurred (defaults to now) document_id: Optional document ID for tracking (always upserts if document already exists) - fact_type_override: Override fact type ('world', 'bank', 'opinion') + fact_type_override: Override fact type ('world', 'interactions', 'opinion') confidence_score: Confidence score for opinions (0.0 to 1.0) Returns: @@ -728,7 +728,7 @@ class MemoryEngine: - "document_id" (optional): Document ID for this specific content item document_id: **DEPRECATED** - Use "document_id" key in each content dict instead. Applies the same document_id to ALL content items that don't specify their own. - fact_type_override: Override fact type for all facts ('world', 'bank', 'opinion') + fact_type_override: Override fact type for all facts ('world', 'interactions', 'opinion') confidence_score: Confidence score for opinions (0.0 to 1.0) Returns: @@ -936,7 +936,7 @@ class MemoryEngine: Args: bank_id: bank ID to recall for query: Recall query - fact_type: List of fact types to recall (e.g., ['world', 'bank']) + fact_type: List of fact types to recall (e.g., ['world', 'interactions']) budget: Budget level for graph traversal (low=100, mid=300, high=600 units) max_tokens: Maximum tokens to return (counts only 'text' field, default 4096) Results are returned until token budget is reached, stopping before @@ -2597,7 +2597,7 @@ Guidelines: logger.info(f"[THINK] Search returned {len(all_results)} results") # Split results by fact type for structured response - agent_results = [r for r in all_results if r.fact_type == 'bank'] + agent_results = [r for r in all_results if r.fact_type == 'interactions'] world_results = [r for r in all_results if r.fact_type == 'world'] opinion_results = [r for r in all_results if r.fact_type == 'opinion'] diff --git a/hindsight-api/hindsight_api/engine/response_models.py b/hindsight-api/hindsight_api/engine/response_models.py index e4dda80c..ee6205b8 100644 --- a/hindsight-api/hindsight_api/engine/response_models.py +++ b/hindsight-api/hindsight_api/engine/response_models.py @@ -61,7 +61,7 @@ class MemoryFact(BaseModel): id: str = Field(description="Unique identifier for the memory fact") text: str = Field(description="The actual text content of the memory") - fact_type: str = Field(description="Type of fact: 'world', 'bank', 'opinion', or 'observation'") + fact_type: str = Field(description="Type of fact: 'world', 'interactions', 'opinion', or 'observation'") entities: Optional[List[str]] = Field(None, description="Entity names mentioned in this fact") context: Optional[str] = Field(None, description="Additional context for the memory") occurred_start: Optional[str] = Field(None, description="ISO format date when the event started occurring") diff --git a/hindsight-api/hindsight_api/engine/retain/fact_extraction.py b/hindsight-api/hindsight_api/engine/retain/fact_extraction.py index 68822a89..4f41f3b0 100644 --- a/hindsight-api/hindsight_api/engine/retain/fact_extraction.py +++ b/hindsight-api/hindsight_api/engine/retain/fact_extraction.py @@ -50,7 +50,7 @@ class Fact(BaseModel): """ # Required fields fact: str = Field(description="Combined fact text: what | when | where | who | why") - fact_type: Literal["world", "bank", "opinion"] = Field(description="Perspective: world/bank/opinion") + fact_type: Literal["world", "interactions", "opinion"] = Field(description="Perspective: world/interactions/opinion") # Optional temporal fields occurred_start: Optional[str] = None @@ -581,20 +581,20 @@ Text: continue # Critical field: fact_type - # LLM uses "assistant" but we convert to "bank" for storage + # LLM uses "assistant" but we convert to "interactions" for storage fact_type = llm_fact.get('fact_type') - # Convert "assistant" → "bank" for storage + # Convert "assistant" → "interactions" for storage if fact_type == 'assistant': - fact_type = 'bank' + fact_type = 'interactions' # Validate fact_type (after conversion) - if fact_type not in ['world', 'bank', 'opinion']: + if fact_type not in ['world', 'interactions', 'opinion']: # Try to fix common mistakes - check if they swapped fact_type and fact_kind fact_kind = llm_fact.get('fact_kind') if fact_kind == 'assistant': - fact_type = 'bank' - elif fact_kind in ['world', 'bank', 'opinion']: + fact_type = 'interactions' + elif fact_kind in ['world', 'interactions', 'opinion']: fact_type = fact_kind else: # Default to 'world' if we can't determine diff --git a/hindsight-api/hindsight_api/engine/retain/link_utils.py b/hindsight-api/hindsight_api/engine/retain/link_utils.py index 45b60c3d..40e8803c 100644 --- a/hindsight-api/hindsight_api/engine/retain/link_utils.py +++ b/hindsight-api/hindsight_api/engine/retain/link_utils.py @@ -529,53 +529,63 @@ async def create_semantic_links_batch( raise -async def insert_entity_links_batch(conn, links: List[tuple], chunk_size: int = 5000): +async def insert_entity_links_batch(conn, links: List[tuple], chunk_size: int = 50000): """ - Insert all entity links in bulk using unnest for efficiency. + Insert all entity links using COPY to temp table + INSERT for maximum speed. - Uses PostgreSQL unnest() to insert many rows in a single query, - which is much faster than executemany over high-latency connections. + Uses PostgreSQL COPY (via copy_records_to_table) for bulk loading, + then INSERT ... ON CONFLICT from temp table. This is the fastest + method for bulk inserts with conflict handling. 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) + chunk_size: Number of rows per batch (default 50000) """ if not links: return 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] + # Create temp table for bulk loading + await conn.execute(""" + CREATE TEMP TABLE IF NOT EXISTS _temp_entity_links ( + from_unit_id uuid, + to_unit_id uuid, + link_type text, + weight float, + entity_id uuid + ) ON COMMIT DROP + """) - # Separate into arrays for unnest - from_ids = [] - to_ids = [] - link_types = [] - weights = [] - entity_ids = [] + # Clear any existing data in temp table + await conn.execute("TRUNCATE _temp_entity_links") - 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 - ) + # Convert links to proper format for COPY + records = [] + for from_id, to_id, link_type, weight, entity_id in links: + records.append(( + uuid_mod.UUID(from_id) if isinstance(from_id, str) else from_id, + uuid_mod.UUID(to_id) if isinstance(to_id, str) else to_id, + link_type, + weight, + 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 - ) + # Bulk load using COPY (fastest method) + await conn.copy_records_to_table( + '_temp_entity_links', + records=records, + columns=['from_unit_id', 'to_unit_id', 'link_type', 'weight', 'entity_id'] + ) + + # Insert from temp table with ON CONFLICT (single query for all rows) + await conn.execute(""" + INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id) + SELECT from_unit_id, to_unit_id, link_type, weight, entity_id + FROM _temp_entity_links + ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING + """) async def create_causal_links_batch( diff --git a/hindsight-api/hindsight_api/engine/retain/types.py b/hindsight-api/hindsight_api/engine/retain/types.py index 1404aa9f..a23f3ee4 100644 --- a/hindsight-api/hindsight_api/engine/retain/types.py +++ b/hindsight-api/hindsight_api/engine/retain/types.py @@ -75,7 +75,7 @@ class ExtractedFact: This is the raw output from fact extraction before processing. """ fact_text: str - fact_type: str # "world", "bank", "opinion", "observation" + fact_type: str # "world", "interactions", "opinion", "observation" entities: List[str] = field(default_factory=list) occurred_start: Optional[datetime] = None occurred_end: Optional[datetime] = None diff --git a/hindsight-api/hindsight_api/models.py b/hindsight-api/hindsight_api/models.py index a8df5924..ab8201bb 100644 --- a/hindsight-api/hindsight_api/models.py +++ b/hindsight-api/hindsight_api/models.py @@ -104,7 +104,7 @@ class MemoryUnit(Base): name="memory_units_document_fkey", ondelete="CASCADE", ), - CheckConstraint("fact_type IN ('world', 'bank', 'opinion', 'observation')"), + CheckConstraint("fact_type IN ('world', 'interactions', 'opinion', 'observation')"), CheckConstraint("confidence_score IS NULL OR (confidence_score >= 0.0 AND confidence_score <= 1.0)"), CheckConstraint( "(fact_type = 'opinion' AND confidence_score IS NOT NULL) OR " diff --git a/uv.lock b/uv.lock index 2bebb93f..b0eece75 100644 --- a/uv.lock +++ b/uv.lock @@ -1141,7 +1141,7 @@ provides-extras = ["test"] [[package]] name = "hindsight-api" -version = "0.0.18" +version = "0.0.17" source = { editable = "hindsight-api" } dependencies = [ { name = "alembic" }, @@ -1243,7 +1243,7 @@ dev = [ [[package]] name = "hindsight-client" -version = "0.0.18" +version = "0.0.17" source = { editable = "hindsight-clients/python" } dependencies = [ { name = "aiohttp" }, @@ -1275,7 +1275,7 @@ provides-extras = ["test"] [[package]] name = "hindsight-dev" -version = "0.0.18" +version = "0.0.17" source = { editable = "hindsight-dev" } dependencies = [ { name = "hindsight-api" },