diff --git a/hindsight-api-slim/hindsight_api/alembic/versions/b3c4d5e6f7a8_add_content_hash_to_chunks.py b/hindsight-api-slim/hindsight_api/alembic/versions/b3c4d5e6f7a8_add_content_hash_to_chunks.py new file mode 100644 index 00000000..d0ec4846 --- /dev/null +++ b/hindsight-api-slim/hindsight_api/alembic/versions/b3c4d5e6f7a8_add_content_hash_to_chunks.py @@ -0,0 +1,32 @@ +"""add content_hash to chunks table for delta retain + +Revision ID: b3c4d5e6f7a8 +Revises: a3b4c5d6e7f8 +Create Date: 2026-03-25 +""" + +from collections.abc import Sequence + +from alembic import context, op + +revision: str = "b3c4d5e6f7a8" +down_revision: str | Sequence[str] | None = "a3b4c5d6e7f8" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def _get_schema_prefix() -> str: + """Get schema prefix for table names (required for multi-tenant support).""" + schema = context.config.get_main_option("target_schema") + return f'"{schema}".' if schema else "" + + +def upgrade() -> None: + schema = _get_schema_prefix() + # Add content_hash column to chunks table for delta comparison + op.execute(f"ALTER TABLE {schema}chunks ADD COLUMN IF NOT EXISTS content_hash TEXT") + + +def downgrade() -> None: + schema = _get_schema_prefix() + op.execute(f"ALTER TABLE {schema}chunks DROP COLUMN IF EXISTS content_hash") diff --git a/hindsight-api-slim/hindsight_api/engine/retain/chunk_storage.py b/hindsight-api-slim/hindsight_api/engine/retain/chunk_storage.py index 34941a10..84ae03ba 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/chunk_storage.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/chunk_storage.py @@ -4,7 +4,9 @@ Chunk storage for retain pipeline. Handles storage of document chunks in the database. """ +import hashlib import logging +from dataclasses import dataclass from ..memory_engine import fq_table from .types import ChunkMetadata @@ -12,6 +14,61 @@ from .types import ChunkMetadata logger = logging.getLogger(__name__) +def compute_chunk_hash(chunk_text: str) -> str: + """Compute SHA256 hash of chunk text for delta comparison.""" + return hashlib.sha256(chunk_text.encode()).hexdigest() + + +@dataclass +class ExistingChunk: + """Represents a chunk already stored in the database.""" + + chunk_id: str + chunk_index: int + content_hash: str | None + + +async def load_existing_chunks(conn, bank_id: str, document_id: str) -> list[ExistingChunk]: + """ + Load existing chunk metadata for a document. + + Returns list of ExistingChunk with chunk_id, chunk_index, and content_hash. + """ + rows = await conn.fetch( + f""" + SELECT chunk_id, chunk_index, content_hash + FROM {fq_table("chunks")} + WHERE document_id = $1 AND bank_id = $2 + ORDER BY chunk_index + """, + document_id, + bank_id, + ) + return [ + ExistingChunk( + chunk_id=row["chunk_id"], + chunk_index=row["chunk_index"], + content_hash=row["content_hash"], + ) + for row in rows + ] + + +async def delete_chunks_by_ids(conn, chunk_ids: list[str]) -> None: + """ + Delete specific chunks by their IDs. + + This cascades to memory_units (via FK with CASCADE delete) + and their links. + """ + if not chunk_ids: + return + await conn.execute( + f"DELETE FROM {fq_table('chunks')} WHERE chunk_id = ANY($1::text[])", + chunk_ids, + ) + + async def store_chunks_batch(conn, bank_id: str, document_id: str, chunks: list[ChunkMetadata]) -> dict[int, str]: """ Store document chunks in the database. @@ -32,6 +89,7 @@ async def store_chunks_batch(conn, bank_id: str, document_id: str, chunks: list[ chunk_ids = [] chunk_texts = [] chunk_indices = [] + content_hashes = [] chunk_id_map = {} for chunk in chunks: @@ -39,19 +97,21 @@ async def store_chunks_batch(conn, bank_id: str, document_id: str, chunks: list[ chunk_ids.append(chunk_id) chunk_texts.append(chunk.chunk_text) chunk_indices.append(chunk.chunk_index) + content_hashes.append(compute_chunk_hash(chunk.chunk_text)) chunk_id_map[chunk.chunk_index] = chunk_id # Batch insert all chunks await conn.execute( f""" - INSERT INTO {fq_table("chunks")} (chunk_id, document_id, bank_id, chunk_text, chunk_index) - SELECT * FROM unnest($1::text[], $2::text[], $3::text[], $4::text[], $5::integer[]) + INSERT INTO {fq_table("chunks")} (chunk_id, document_id, bank_id, chunk_text, chunk_index, content_hash) + SELECT * FROM unnest($1::text[], $2::text[], $3::text[], $4::text[], $5::integer[], $6::text[]) """, chunk_ids, [document_id] * len(chunk_texts), [bank_id] * len(chunk_texts), chunk_texts, chunk_indices, + content_hashes, ) return chunk_id_map diff --git a/hindsight-api-slim/hindsight_api/engine/retain/fact_storage.py b/hindsight-api-slim/hindsight_api/engine/retain/fact_storage.py index 412ecbfb..43b9e93c 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/fact_storage.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/fact_storage.py @@ -221,7 +221,10 @@ async def handle_document_tracking( document_tags: list[str] | None = None, ) -> None: """ - Handle document tracking in the database. + Handle document tracking in the database (full-replace mode). + + Deletes the existing document (cascading to all units and links) on the + first batch, then inserts the new document record. Args: conn: Database connection @@ -238,14 +241,51 @@ async def handle_document_tracking( combined_content = _sanitize_text(combined_content) or "" content_hash = hashlib.sha256(combined_content.encode()).hexdigest() - # Always delete old document first if it exists (cascades to units and links) + # Delete old document first (cascades to units and links) # Only delete on the first batch to avoid deleting data we just inserted if is_first_batch: await conn.fetchval( - f"DELETE FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 RETURNING id", document_id, bank_id + f"DELETE FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 RETURNING id", + document_id, + bank_id, ) # Insert document (or update if exists from concurrent operations) + await _upsert_document_row(conn, bank_id, document_id, combined_content, content_hash, retain_params, document_tags) + + +async def upsert_document_metadata( + conn, + bank_id: str, + document_id: str, + combined_content: str, + retain_params: dict | None = None, + document_tags: list[str] | None = None, +) -> None: + """ + Update document metadata without deleting existing facts/chunks. + + Used by delta retain: the document row is upserted but chunks and + memory_units are managed separately at the chunk level. + """ + import hashlib + + combined_content = _sanitize_text(combined_content) or "" + content_hash = hashlib.sha256(combined_content.encode()).hexdigest() + + await _upsert_document_row(conn, bank_id, document_id, combined_content, content_hash, retain_params, document_tags) + + +async def _upsert_document_row( + conn, + bank_id: str, + document_id: str, + combined_content: str, + content_hash: str, + retain_params: dict | None = None, + document_tags: list[str] | None = None, +) -> None: + """Insert or update a document row.""" await conn.execute( f""" INSERT INTO {fq_table("documents")} (id, bank_id, original_text, content_hash, metadata, retain_params, tags) @@ -266,3 +306,34 @@ async def handle_document_tracking( json.dumps(retain_params) if retain_params else None, document_tags or [], ) + + +async def update_memory_units_tags( + conn, + bank_id: str, + document_id: str, + tags: list[str], +) -> int: + """ + Update tags on all memory_units belonging to a document. + + Used during delta retain to propagate tag changes to unchanged facts. + + Returns: + Number of memory units updated. + """ + result = await conn.execute( + f""" + UPDATE {fq_table("memory_units")} + SET tags = $3, updated_at = NOW() + WHERE bank_id = $1 AND document_id = $2 + """, + bank_id, + document_id, + tags or [], + ) + # result is a status string like "UPDATE 5" + try: + return int(result.split()[-1]) + except (ValueError, IndexError): + return 0 diff --git a/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py b/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py index f5c0cd3a..54a8d735 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py @@ -7,6 +7,7 @@ Coordinates all retain pipeline modules to store memories efficiently. import logging import time import uuid +from collections import defaultdict from collections.abc import Awaitable, Callable from datetime import UTC, datetime from typing import Any @@ -64,11 +65,165 @@ from . import ( fact_storage, link_creation, ) -from .types import EntityLink, ExtractedFact, ProcessedFact, RetainContent, RetainContentDict +from .types import ChunkMetadata, EntityLink, ExtractedFact, ProcessedFact, RetainContent, RetainContentDict logger = logging.getLogger(__name__) +def _build_retain_params(contents_dicts, document_tags=None, doc_contents=None): + """Build retain_params and merged_tags from content dicts.""" + if doc_contents is not None: + # Per-document mode: doc_contents is list of (idx, content_dict) + items = [item for _, item in doc_contents] + else: + items = contents_dicts + + all_tags = set(document_tags or []) + for item in items: + item_tags = item.get("tags", []) or [] + all_tags.update(item_tags) + merged_tags = list(all_tags) + + retain_params = {} + if items: + first_item = items[0] + if first_item.get("context"): + retain_params["context"] = first_item["context"] + if first_item.get("event_date"): + retain_params["event_date"] = ( + first_item["event_date"].isoformat() + if hasattr(first_item["event_date"], "isoformat") + else str(first_item["event_date"]) + ) + if first_item.get("metadata"): + retain_params["metadata"] = first_item["metadata"] + + return retain_params, merged_tags + + +async def _insert_facts_and_links( + conn, + entity_resolver, + bank_id: str, + contents: list[RetainContent], + extracted_facts: list, + processed_facts: list[ProcessedFact], + config, + log_buffer: list[str], + outbox_callback=None, +) -> list[list[str]]: + """ + Shared pipeline: insert facts, process entities, create all link types. + + Used by both the full retain and delta retain paths. + + Returns: + List of unit ID lists mapped back to original content items. + """ + unit_ids = await fact_storage.insert_facts_batch(conn, bank_id, processed_facts) + step_start = time.time() + log_buffer.append(f" Insert facts: {len(unit_ids)} units in {time.time() - step_start:.3f}s") + + if unit_ids: + # Process entities + step_start = time.time() + user_entities_per_content = { + idx: content.entities for idx, content in enumerate(contents) if content.entities + } + entity_links = await entity_processing.process_entities_batch( + entity_resolver, + conn, + bank_id, + unit_ids, + processed_facts, + log_buffer, + user_entities_per_content=user_entities_per_content, + entity_labels=getattr(config, "entity_labels", None), + ) + log_buffer.append(f" Process entities: {len(entity_links)} links in {time.time() - step_start:.3f}s") + + # Create temporal links + step_start = time.time() + temporal_link_count = await link_creation.create_temporal_links_batch(conn, bank_id, unit_ids) + log_buffer.append(f" Temporal links: {temporal_link_count} links in {time.time() - step_start:.3f}s") + + # Create semantic links + step_start = time.time() + embeddings_for_links = [fact.embedding for fact in processed_facts] + semantic_link_count = await link_creation.create_semantic_links_batch( + conn, bank_id, unit_ids, embeddings_for_links + ) + log_buffer.append(f" Semantic links: {semantic_link_count} links in {time.time() - step_start:.3f}s") + + # Insert entity links + step_start = time.time() + if entity_links: + await entity_processing.insert_entity_links_batch(conn, entity_links) + log_buffer.append( + f" Entity links: {len(entity_links) if entity_links else 0} links in {time.time() - step_start:.3f}s" + ) + + # Create causal links + step_start = time.time() + causal_link_count = await link_creation.create_causal_links_batch(conn, unit_ids, processed_facts) + log_buffer.append(f" Causal links: {causal_link_count} links in {time.time() - step_start:.3f}s") + + # Map results back to original content items + result_unit_ids = _map_results_to_contents(contents, extracted_facts, unit_ids if unit_ids else []) + + if outbox_callback: + await outbox_callback(conn) + + return result_unit_ids + + +async def _extract_and_embed( + contents: list[RetainContent], + llm_config, + agent_name: str, + config, + embeddings_model, + format_date_fn, + fact_type_override: str | None, + log_buffer: list[str], + pool=None, + operation_id: str | None = None, + schema: str | None = None, +) -> tuple[list, list[ProcessedFact], list[ChunkMetadata], TokenUsage]: + """ + Shared pipeline: extract facts from contents and generate embeddings. + + Returns: + Tuple of (extracted_facts, processed_facts, chunks_metadata, usage) + """ + step_start = time.time() + extracted_facts, chunks, usage = await fact_extraction.extract_facts_from_contents( + contents, llm_config, agent_name, config, pool, operation_id, schema + ) + log_buffer.append( + f" Extract facts: {len(extracted_facts)} facts, {len(chunks)} chunks " + f"from {len(contents)} contents in {time.time() - step_start:.3f}s" + ) + + if not extracted_facts: + return extracted_facts, [], chunks, usage + + if fact_type_override: + for fact in extracted_facts: + fact.fact_type = fact_type_override + + step_start = time.time() + augmented_texts = embedding_processing.augment_texts_with_dates(extracted_facts, format_date_fn) + embeddings = await embedding_processing.generate_embeddings_batch(embeddings_model, augmented_texts) + log_buffer.append(f" Generate embeddings: {len(embeddings)} embeddings in {time.time() - step_start:.3f}s") + + processed_facts = [ + ProcessedFact.from_extracted_fact(ef, emb) for ef, emb in zip(extracted_facts, embeddings) + ] + + return extracted_facts, processed_facts, chunks, usage + + async def retain_batch( pool, embeddings_model, @@ -90,28 +245,13 @@ async def retain_batch( """ Process a batch of content through the retain pipeline. - Args: - pool: Database connection pool - embeddings_model: Embeddings model for generating embeddings - llm_config: LLM configuration for fact extraction - entity_resolver: Entity resolver for entity processing - format_date_fn: Function to format datetime to readable string - bank_id: Bank identifier - contents_dicts: List of content dictionaries - config: Resolved HindsightConfig for this bank - document_id: Optional document ID - is_first_batch: Whether this is the first batch - fact_type_override: Override fact type for all facts - confidence_score: Confidence score for opinions - document_tags: Tags applied to all items in this batch - - Returns: - Tuple of (unit ID lists, token usage for fact extraction) + Supports delta retain: when upserting a document that already has chunks, + only re-processes chunks whose content has changed. Unchanged chunks keep + their existing facts, entities, and links. """ start_time = time.time() total_chars = sum(len(item.get("content", "")) for item in contents_dicts) - # Buffer all logs log_buffer = [] log_buffer.append(f"{'=' * 60}") log_buffer.append(f"RETAIN_BATCH START: {bank_id}") @@ -123,20 +263,373 @@ async def retain_batch( agent_name = profile["name"] # Convert dicts to RetainContent objects + contents = _build_contents(contents_dicts, document_tags) + + # --- Delta retain: check if we can skip unchanged chunks --- + if is_first_batch: + delta_result = await _try_delta_retain( + pool, embeddings_model, llm_config, entity_resolver, format_date_fn, + bank_id, contents_dicts, contents, config, document_id, fact_type_override, + document_tags, agent_name, log_buffer, start_time, operation_id, schema, outbox_callback, + ) + if delta_result is not None: + return delta_result + + # --- Full retain path --- + extracted_facts, processed_facts, chunks, usage = await _extract_and_embed( + contents, llm_config, agent_name, config, embeddings_model, format_date_fn, + fact_type_override, log_buffer, pool, operation_id, schema, + ) + + if not extracted_facts: + await _handle_zero_facts_documents( + pool, bank_id, contents_dicts, contents, config, document_id, + is_first_batch, document_tags, chunks, log_buffer, start_time, + ) + return [[] for _ in contents], usage + + # Group contents by document_id + contents_by_doc = defaultdict(list) + for idx, content_dict in enumerate(contents_dicts): + doc_id = content_dict.get("document_id") + contents_by_doc[doc_id].append((idx, content_dict)) + + # Database transaction (retried on deadlock) + result_unit_ids: list[list[str]] = [] + log_buffer_pre_db = len(log_buffer) + + async def _run_db_work() -> None: + nonlocal result_unit_ids + del log_buffer[log_buffer_pre_db:] + document_ids_added: list[str] = [] + for pf in processed_facts: + pf.document_id = None + pf.chunk_id = None + entity_resolver.discard_pending_stats() + + async with acquire_with_retry(pool) as conn: + async with conn.transaction(): + # Handle document tracking + step_start = time.time() + doc_id_mapping = {} + + if document_id: + combined_content = "\n".join([c.get("content", "") for c in contents_dicts]) + retain_params, merged_tags = _build_retain_params(contents_dicts, document_tags) + await fact_storage.handle_document_tracking( + conn, bank_id, document_id, combined_content, is_first_batch, retain_params, merged_tags + ) + document_ids_added.append(document_id) + doc_id_mapping[None] = document_id + else: + has_any_doc_ids = any(item.get("document_id") for item in contents_dicts) + if has_any_doc_ids or chunks: + for original_doc_id, doc_contents in contents_by_doc.items(): + actual_doc_id = original_doc_id + should_create_doc = (original_doc_id is not None) or chunks + if should_create_doc: + if actual_doc_id is None: + actual_doc_id = str(uuid.uuid4()) + doc_id_mapping[original_doc_id] = actual_doc_id + combined_content = "\n".join([c.get("content", "") for _, c in doc_contents]) + retain_params, merged_tags = _build_retain_params( + contents_dicts, document_tags, doc_contents=doc_contents + ) + await fact_storage.handle_document_tracking( + conn, bank_id, actual_doc_id, combined_content, + is_first_batch, retain_params, merged_tags, + ) + document_ids_added.append(actual_doc_id) + + if document_ids_added: + log_buffer.append( + f" Document tracking: {len(document_ids_added)} documents in {time.time() - step_start:.3f}s" + ) + + # Store chunks and map to facts + step_start = time.time() + chunk_id_map_by_doc = {} + if chunks: + chunks_by_doc = defaultdict(list) + for chunk in chunks: + original_doc_id = contents_dicts[chunk.content_index].get("document_id") + actual_doc_id = doc_id_mapping.get(original_doc_id, original_doc_id) + if actual_doc_id is None and document_id: + actual_doc_id = document_id + chunks_by_doc[actual_doc_id].append(chunk) + + for doc_id, doc_chunks in chunks_by_doc.items(): + chunk_id_map = await chunk_storage.store_chunks_batch(conn, bank_id, doc_id, doc_chunks) + for chunk_idx, chunk_id in chunk_id_map.items(): + chunk_id_map_by_doc[(doc_id, chunk_idx)] = chunk_id + + log_buffer.append( + f" Store chunks: {len(chunks)} chunks for {len(chunks_by_doc)} documents " + f"in {time.time() - step_start:.3f}s" + ) + + # Map chunk_ids and document_ids to facts + for fact, processed_fact in zip(extracted_facts, processed_facts): + original_doc_id = contents_dicts[fact.content_index].get("document_id") + actual_doc_id = doc_id_mapping.get(original_doc_id, original_doc_id) + if actual_doc_id is None and document_id: + actual_doc_id = document_id + processed_fact.document_id = actual_doc_id + if chunks and fact.chunk_index is not None: + chunk_id = chunk_id_map_by_doc.get((actual_doc_id, fact.chunk_index)) + if chunk_id: + processed_fact.chunk_id = chunk_id + + # Insert facts and create all links (shared pipeline) + result_unit_ids = await _insert_facts_and_links( + conn, entity_resolver, bank_id, contents, extracted_facts, + processed_facts, config, log_buffer, outbox_callback, + ) + + await entity_resolver.flush_pending_stats() + + total_time = time.time() - start_time + log_buffer.append(f"{'=' * 60}") + log_buffer.append(f"RETAIN_BATCH COMPLETE: {len(processed_facts)} units in {total_time:.3f}s") + if document_ids_added: + log_buffer.append(f"Documents: {', '.join(document_ids_added)}") + log_buffer.append(f"{'=' * 60}") + logger.info("\n" + "\n".join(log_buffer) + "\n") + + await retry_with_backoff(_run_db_work) + return result_unit_ids, usage + + +# --------------------------------------------------------------------------- +# Delta retain +# --------------------------------------------------------------------------- + + +async def _try_delta_retain( + pool, embeddings_model, llm_config, entity_resolver, format_date_fn, + bank_id, contents_dicts, contents, config, document_id, fact_type_override, + document_tags, agent_name, log_buffer, start_time, operation_id, schema, outbox_callback, +): + """ + Attempt delta retain for a document upsert. Returns result tuple if delta + was performed, or None to fall back to full retain. + """ + # Need a single document_id + effective_doc_id = document_id + if not effective_doc_id: + doc_ids = {item.get("document_id") for item in contents_dicts if item.get("document_id")} + if len(doc_ids) != 1: + return None + effective_doc_id = doc_ids.pop() + + # Load existing chunks + async with acquire_with_retry(pool) as conn: + existing_chunks = await chunk_storage.load_existing_chunks(conn, bank_id, effective_doc_id) + + if not existing_chunks: + return None + + if any(c.content_hash is None for c in existing_chunks): + logger.info(f"Delta retain skipped for {effective_doc_id}: existing chunks lack content_hash (pre-migration)") + return None + + # Chunk new content and classify changes + step_start = time.time() + new_chunks_with_contents = _chunk_contents_for_delta(contents, config) + log_buffer.append( + f"[delta] Chunked new content: {len(new_chunks_with_contents)} chunks in {time.time() - step_start:.3f}s" + ) + + existing_by_index = {c.chunk_index: c for c in existing_chunks} + new_hashes = {idx: chunk_storage.compute_chunk_hash(text) for idx, text in new_chunks_with_contents.items()} + + unchanged_indices, changed_indices, new_indices, removed_indices = [], [], [], [] + for idx, new_hash in new_hashes.items(): + existing = existing_by_index.get(idx) + if existing and existing.content_hash == new_hash: + unchanged_indices.append(idx) + elif existing: + changed_indices.append(idx) + else: + new_indices.append(idx) + for idx in existing_by_index: + if idx not in new_hashes: + removed_indices.append(idx) + + log_buffer.append( + f"[delta] Chunk diff: {len(unchanged_indices)} unchanged, " + f"{len(changed_indices)} changed, {len(new_indices)} new, " + f"{len(removed_indices)} removed" + ) + + if not unchanged_indices: + logger.info(f"Delta retain: no unchanged chunks for {effective_doc_id}, falling back to full retain") + return None + + chunks_to_process = changed_indices + new_indices + + if not chunks_to_process and not removed_indices: + # Nothing changed — just update document metadata/tags + log_buffer.append("[delta] No chunk changes detected — updating document metadata only") + return await _delta_metadata_only( + pool, bank_id, contents_dicts, contents, effective_doc_id, + document_tags, log_buffer, start_time, outbox_callback, + ) + + # Build content items for only the changed/new chunks + delta_contents, delta_chunk_map = _build_delta_contents(contents, new_chunks_with_contents, chunks_to_process) + + if not delta_contents: + return await _delta_metadata_only( + pool, bank_id, contents_dicts, contents, effective_doc_id, + document_tags, log_buffer, start_time, outbox_callback, + ) + + # Extract facts and generate embeddings (shared pipeline) + extracted_facts, processed_facts, new_chunk_metadata, usage = await _extract_and_embed( + delta_contents, llm_config, agent_name, config, embeddings_model, format_date_fn, + fact_type_override, log_buffer, pool, operation_id, schema, + ) + + # Database transaction + result_unit_ids: list[list[str]] = [] + log_buffer_pre_db = len(log_buffer) + + async def _run_delta_db_work() -> None: + nonlocal result_unit_ids + del log_buffer[log_buffer_pre_db:] + for pf in processed_facts: + pf.document_id = None + pf.chunk_id = None + entity_resolver.discard_pending_stats() + + async with acquire_with_retry(pool) as conn: + async with conn.transaction(): + # Update document metadata (no delete) + step_start = time.time() + combined_content = "\n".join([c.get("content", "") for c in contents_dicts]) + retain_params, merged_tags = _build_retain_params(contents_dicts, document_tags) + await fact_storage.upsert_document_metadata( + conn, bank_id, effective_doc_id, combined_content, retain_params, merged_tags, + ) + log_buffer.append(f" Document metadata update in {time.time() - step_start:.3f}s") + + # Delete changed and removed chunks (cascades to memory_units and links) + step_start = time.time() + chunks_to_delete = [ + existing_by_index[idx].chunk_id + for idx in changed_indices + removed_indices + if idx in existing_by_index + ] + await chunk_storage.delete_chunks_by_ids(conn, chunks_to_delete) + log_buffer.append( + f" Deleted {len(chunks_to_delete)} chunks " + f"({len(changed_indices)} changed + {len(removed_indices)} removed) " + f"in {time.time() - step_start:.3f}s" + ) + + # Update tags on unchanged chunks' memory units + step_start = time.time() + updated_count = await fact_storage.update_memory_units_tags( + conn, bank_id, effective_doc_id, merged_tags + ) + log_buffer.append( + f" Updated tags on {updated_count} existing memory units in {time.time() - step_start:.3f}s" + ) + + # Store new/changed chunks + step_start = time.time() + chunk_id_map_by_doc = {} + if new_chunk_metadata: + remapped_chunks = [ + ChunkMetadata( + chunk_text=cm.chunk_text, + fact_count=cm.fact_count, + content_index=cm.content_index, + chunk_index=delta_chunk_map.get(cm.chunk_index, cm.chunk_index), + ) + for cm in new_chunk_metadata + ] + chunk_id_map = await chunk_storage.store_chunks_batch( + conn, bank_id, effective_doc_id, remapped_chunks + ) + for chunk_idx, chunk_id in chunk_id_map.items(): + chunk_id_map_by_doc[(effective_doc_id, chunk_idx)] = chunk_id + log_buffer.append( + f" Stored {len(remapped_chunks)} new/changed chunks in {time.time() - step_start:.3f}s" + ) + + # Map chunk_ids and document_ids to processed facts + for ef, pf in zip(extracted_facts, processed_facts): + pf.document_id = effective_doc_id + if ef.chunk_index is not None: + original_idx = delta_chunk_map.get(ef.chunk_index, ef.chunk_index) + chunk_id = chunk_id_map_by_doc.get((effective_doc_id, original_idx)) + if chunk_id: + pf.chunk_id = chunk_id + + # Insert facts and create all links (shared pipeline) + result_unit_ids = await _insert_facts_and_links( + conn, entity_resolver, bank_id, contents, extracted_facts, + processed_facts, config, log_buffer, outbox_callback, + ) + + await entity_resolver.flush_pending_stats() + + total_time = time.time() - start_time + log_buffer.append(f"{'=' * 60}") + log_buffer.append( + f"DELTA RETAIN COMPLETE: {len(processed_facts)} new units, " + f"{len(unchanged_indices)} chunks unchanged in {total_time:.3f}s" + ) + log_buffer.append(f"Document: {effective_doc_id}") + log_buffer.append(f"{'=' * 60}") + logger.info("\n" + "\n".join(log_buffer) + "\n") + + await retry_with_backoff(_run_delta_db_work) + return result_unit_ids, usage + + +async def _delta_metadata_only( + pool, bank_id, contents_dicts, contents, document_id, document_tags, + log_buffer, start_time, outbox_callback, +): + """Handle the case where no chunks changed — just update document metadata and tags.""" + async with acquire_with_retry(pool) as conn: + async with conn.transaction(): + combined_content = "\n".join([c.get("content", "") for c in contents_dicts]) + retain_params, merged_tags = _build_retain_params(contents_dicts, document_tags) + await fact_storage.upsert_document_metadata( + conn, bank_id, document_id, combined_content, retain_params, merged_tags, + ) + await fact_storage.update_memory_units_tags(conn, bank_id, document_id, merged_tags) + if outbox_callback: + await outbox_callback(conn) + + total_time = time.time() - start_time + log_buffer.append(f"DELTA RETAIN (no changes): metadata updated in {total_time:.3f}s") + logger.info("\n" + "\n".join(log_buffer) + "\n") + return [[] for _ in contents], TokenUsage() + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _build_contents(contents_dicts: list[RetainContentDict], document_tags: list[str] | None) -> list[RetainContent]: + """Convert content dicts to RetainContent objects.""" contents = [] for item in contents_dicts: - # Merge item-level tags with document-level tags item_tags = item.get("tags", []) or [] merged_tags = list(set(item_tags + (document_tags or []))) - # Handle event_date: distinguish "not provided" (default to now) from - # "explicitly None" (caller opted into no timestamp). if "event_date" in item and item["event_date"] is None: - event_date_value = None # Caller explicitly signalled "unknown date" + event_date_value = None elif item.get("event_date"): event_date_value = parse_datetime_flexible(item["event_date"]) else: - event_date_value = utcnow() # Backward-compatible default + event_date_value = utcnow() content = RetainContent( content=item["content"], @@ -148,382 +641,107 @@ async def retain_batch( observation_scopes=item.get("observation_scopes"), ) contents.append(content) + return contents - # Step 1: Extract facts from all contents - step_start = time.time() - extracted_facts, chunks, usage = await fact_extraction.extract_facts_from_contents( - contents, llm_config, agent_name, config, pool, operation_id, schema - ) - log_buffer.append( - f"[1] Extract facts: {len(extracted_facts)} facts, {len(chunks)} chunks from {len(contents)} contents in {time.time() - step_start:.3f}s" +async def _handle_zero_facts_documents( + pool, bank_id, contents_dicts, contents, config, document_id, + is_first_batch, document_tags, chunks, log_buffer, start_time, +): + """Handle document tracking when zero facts were extracted.""" + docs_tracked = 0 + async with acquire_with_retry(pool) as conn: + async with conn.transaction(): + contents_by_doc = defaultdict(list) + for idx, content_dict in enumerate(contents_dicts): + doc_id = content_dict.get("document_id") + contents_by_doc[doc_id].append((idx, content_dict)) + + if document_id: + combined_content = "\n".join([c.get("content", "") for c in contents_dicts]) + retain_params, merged_tags = _build_retain_params(contents_dicts, document_tags) + await fact_storage.handle_document_tracking( + conn, bank_id, document_id, combined_content, is_first_batch, retain_params, merged_tags + ) + docs_tracked += 1 + else: + has_any_doc_ids = any(item.get("document_id") for item in contents_dicts) + if has_any_doc_ids or chunks: + for original_doc_id, doc_contents in contents_by_doc.items(): + should_create_doc = (original_doc_id is not None) or chunks + if not should_create_doc: + continue + actual_doc_id = original_doc_id or str(uuid.uuid4()) + combined_content = "\n".join([c.get("content", "") for _, c in doc_contents]) + retain_params, merged_tags = _build_retain_params( + contents_dicts, document_tags, doc_contents=doc_contents + ) + await fact_storage.handle_document_tracking( + conn, bank_id, actual_doc_id, combined_content, + is_first_batch, retain_params, merged_tags, + ) + docs_tracked += 1 + + total_time = time.time() - start_time + doc_status = f"{docs_tracked} document(s) tracked" if docs_tracked > 0 else "no document tracked" + logger.info( + f"RETAIN_BATCH COMPLETE: 0 facts extracted from {len(contents)} contents " + f"in {total_time:.3f}s ({doc_status}, no facts)" ) - if not extracted_facts: - # Still need to create document if document_id was provided or chunks exist - from collections import defaultdict - docs_tracked = 0 - async with acquire_with_retry(pool) as conn: - async with conn.transaction(): - # Group contents by document_id (consistent with normal path) - contents_by_doc_early = defaultdict(list) - for idx, content_dict in enumerate(contents_dicts): - doc_id = content_dict.get("document_id") - contents_by_doc_early[doc_id].append((idx, content_dict)) +def _chunk_contents_for_delta(contents: list[RetainContent], config) -> dict[int, str]: + """ + Chunk contents the same way fact_extraction does, returning a map of + global_chunk_index -> chunk_text. + """ + result = {} + global_chunk_idx = 0 + for content in contents: + chunk_size = getattr(config, "retain_chunk_size", 120000) + chunks = fact_extraction.chunk_text(content.content, chunk_size) + for chunk_text in chunks: + result[global_chunk_idx] = chunk_text + global_chunk_idx += 1 + return result - if document_id: - # Legacy: single document_id parameter - combined_content = "\n".join([c.get("content", "") for c in contents_dicts]) - # Collect tags from all content items and merge with document_tags - all_tags = set(document_tags or []) - for item in contents_dicts: - item_tags = item.get("tags", []) or [] - all_tags.update(item_tags) - merged_tags = list(all_tags) - retain_params = {} - if contents_dicts: - first_item = contents_dicts[0] - if first_item.get("context"): - retain_params["context"] = first_item["context"] - if first_item.get("event_date"): - retain_params["event_date"] = ( - first_item["event_date"].isoformat() - if hasattr(first_item["event_date"], "isoformat") - else str(first_item["event_date"]) - ) - if first_item.get("metadata"): - retain_params["metadata"] = first_item["metadata"] - await fact_storage.handle_document_tracking( - conn, bank_id, document_id, combined_content, is_first_batch, retain_params, merged_tags - ) - docs_tracked += 1 - else: - # Handle per-item document_ids and/or chunks (mirrors normal path logic) - has_any_doc_ids = any(item.get("document_id") for item in contents_dicts) +def _build_delta_contents( + original_contents: list[RetainContent], + new_chunks_with_contents: dict[int, str], + chunks_to_process: list[int], +) -> tuple[list[RetainContent], dict[int, int]]: + """ + Build RetainContent items containing only the chunks that need processing. - if has_any_doc_ids or chunks: - for original_doc_id, doc_contents in contents_by_doc_early.items(): - should_create_doc = (original_doc_id is not None) or chunks - if not should_create_doc: - continue + Returns: + - List of RetainContent items (one per chunk to process) + - Map of delta_chunk_index -> original_chunk_index + """ + if not chunks_to_process or not original_contents: + return [], {} - actual_doc_id = original_doc_id - if actual_doc_id is None: - # No document_id but have chunks - generate one - actual_doc_id = str(uuid.uuid4()) + template_content = original_contents[0] + delta_contents = [] + delta_chunk_map = {} - combined_content = "\n".join([c.get("content", "") for _, c in doc_contents]) - all_tags = set(document_tags or []) - for _, item in doc_contents: - item_tags = item.get("tags", []) or [] - all_tags.update(item_tags) - merged_tags = list(all_tags) - - retain_params = {} - if doc_contents: - first_item = doc_contents[0][1] - if first_item.get("context"): - retain_params["context"] = first_item["context"] - if first_item.get("event_date"): - retain_params["event_date"] = ( - first_item["event_date"].isoformat() - if hasattr(first_item["event_date"], "isoformat") - else str(first_item["event_date"]) - ) - if first_item.get("metadata"): - retain_params["metadata"] = first_item["metadata"] - await fact_storage.handle_document_tracking( - conn, - bank_id, - actual_doc_id, - combined_content, - is_first_batch, - retain_params, - merged_tags, - ) - docs_tracked += 1 - - total_time = time.time() - start_time - doc_status = f"{docs_tracked} document(s) tracked" if docs_tracked > 0 else "no document tracked" - logger.info( - f"RETAIN_BATCH COMPLETE: 0 facts extracted from {len(contents)} contents in {total_time:.3f}s ({doc_status}, no facts)" + for original_chunk_idx in sorted(chunks_to_process): + chunk_text = new_chunks_with_contents.get(original_chunk_idx) + if not chunk_text: + continue + delta_content = RetainContent( + content=chunk_text, + context=template_content.context, + event_date=template_content.event_date, + metadata=template_content.metadata, + entities=template_content.entities, + tags=template_content.tags, + observation_scopes=template_content.observation_scopes, ) - return [[] for _ in contents], usage + delta_contents.append(delta_content) + delta_chunk_map[len(delta_contents) - 1] = original_chunk_idx - # Apply fact_type_override if provided - if fact_type_override: - for fact in extracted_facts: - fact.fact_type = fact_type_override - - # Step 2: Augment texts and generate embeddings - step_start = time.time() - augmented_texts = embedding_processing.augment_texts_with_dates(extracted_facts, format_date_fn) - embeddings = await embedding_processing.generate_embeddings_batch(embeddings_model, augmented_texts) - log_buffer.append(f"[2] Generate embeddings: {len(embeddings)} embeddings in {time.time() - step_start:.3f}s") - - # Step 3: Convert to ProcessedFact objects (without chunk_ids yet) - processed_facts = [ - ProcessedFact.from_extracted_fact(extracted_fact, embedding) - for extracted_fact, embedding in zip(extracted_facts, embeddings) - ] - - # Group contents by document_id for document tracking and chunk storage - from collections import defaultdict - - contents_by_doc = defaultdict(list) - for idx, content_dict in enumerate(contents_dicts): - doc_id = content_dict.get("document_id") - contents_by_doc[doc_id].append((idx, content_dict)) - - # Step 4: Database transaction (retried on deadlock) - result_unit_ids: list[list[str]] = [] - - log_buffer_pre_db = len(log_buffer) - - async def _run_db_work() -> None: - nonlocal result_unit_ids - - # Reset per-fact mutations and log buffer so each retry attempt starts clean - del log_buffer[log_buffer_pre_db:] - document_ids_added: list[str] = [] - for pf in processed_facts: - pf.document_id = None - pf.chunk_id = None - - # Discard any leftover pending stats from a previous failed attempt so - # retries don't double-count or accumulate unbounded state. - entity_resolver.discard_pending_stats() - - async with acquire_with_retry(pool) as conn: - async with conn.transaction(): - # Handle document tracking for all documents - step_start = time.time() - # Map None document_id to generated UUIDs - doc_id_mapping = {} # Maps original doc_id (including None) to actual doc_id used - - if document_id: - # Legacy: single document_id parameter - combined_content = "\n".join([c.get("content", "") for c in contents_dicts]) - retain_params = {} - # Collect tags from all content items and merge with document_tags - all_tags = set(document_tags or []) - for item in contents_dicts: - item_tags = item.get("tags", []) or [] - all_tags.update(item_tags) - merged_tags = list(all_tags) - - if contents_dicts: - first_item = contents_dicts[0] - if first_item.get("context"): - retain_params["context"] = first_item["context"] - if first_item.get("event_date"): - retain_params["event_date"] = ( - first_item["event_date"].isoformat() - if hasattr(first_item["event_date"], "isoformat") - else str(first_item["event_date"]) - ) - if first_item.get("metadata"): - retain_params["metadata"] = first_item["metadata"] - - await fact_storage.handle_document_tracking( - conn, bank_id, document_id, combined_content, is_first_batch, retain_params, merged_tags - ) - document_ids_added.append(document_id) - doc_id_mapping[None] = document_id # For backwards compatibility - else: - # Handle per-item document_ids (create documents if any item has document_id or if chunks exist) - has_any_doc_ids = any(item.get("document_id") for item in contents_dicts) - - if has_any_doc_ids or chunks: - for original_doc_id, doc_contents in contents_by_doc.items(): - actual_doc_id = original_doc_id - - # Only create document record if: - # 1. Item has explicit document_id, OR - # 2. There are chunks (need document for chunk storage) - should_create_doc = (original_doc_id is not None) or chunks - - if should_create_doc: - if actual_doc_id is None: - # No document_id but have chunks - generate one - actual_doc_id = str(uuid.uuid4()) - - # Store mapping for later use - doc_id_mapping[original_doc_id] = actual_doc_id - - # Combine content for this document - combined_content = "\n".join([c.get("content", "") for _, c in doc_contents]) - - # Collect tags from all content items for this document and merge with document_tags - all_tags = set(document_tags or []) - for _, item in doc_contents: - item_tags = item.get("tags", []) or [] - all_tags.update(item_tags) - merged_tags = list(all_tags) - - # Extract retain params from first content item - retain_params = {} - if doc_contents: - first_item = doc_contents[0][1] - if first_item.get("context"): - retain_params["context"] = first_item["context"] - if first_item.get("event_date"): - retain_params["event_date"] = ( - first_item["event_date"].isoformat() - if hasattr(first_item["event_date"], "isoformat") - else str(first_item["event_date"]) - ) - if first_item.get("metadata"): - retain_params["metadata"] = first_item["metadata"] - - await fact_storage.handle_document_tracking( - conn, - bank_id, - actual_doc_id, - combined_content, - is_first_batch, - retain_params, - merged_tags, - ) - document_ids_added.append(actual_doc_id) - - if document_ids_added: - log_buffer.append( - f"[2.5] Document tracking: {len(document_ids_added)} documents in {time.time() - step_start:.3f}s" - ) - - # Store chunks and map to facts for all documents - step_start = time.time() - chunk_id_map_by_doc = {} # Maps (doc_id, chunk_index) -> chunk_id - - if chunks: - # Group chunks by their source document - chunks_by_doc = defaultdict(list) - for chunk in chunks: - # chunk.content_index tells us which content this chunk came from - original_doc_id = contents_dicts[chunk.content_index].get("document_id") - # Map to actual document_id (handles None -> generated UUID mapping) - actual_doc_id = doc_id_mapping.get(original_doc_id, original_doc_id) - if actual_doc_id is None and document_id: - actual_doc_id = document_id - chunks_by_doc[actual_doc_id].append(chunk) - - # Store chunks for each document - for doc_id, doc_chunks in chunks_by_doc.items(): - chunk_id_map = await chunk_storage.store_chunks_batch(conn, bank_id, doc_id, doc_chunks) - # Store mapping with document context - for chunk_idx, chunk_id in chunk_id_map.items(): - chunk_id_map_by_doc[(doc_id, chunk_idx)] = chunk_id - - log_buffer.append( - f"[3] Store chunks: {len(chunks)} chunks for {len(chunks_by_doc)} documents in {time.time() - step_start:.3f}s" - ) - - # Map chunk_ids and document_ids to facts - for fact, processed_fact in zip(extracted_facts, processed_facts): - # Get the original document_id for this fact's source content - original_doc_id = contents_dicts[fact.content_index].get("document_id") - # Map to actual document_id (handles None -> generated UUID mapping) - actual_doc_id = doc_id_mapping.get(original_doc_id, original_doc_id) - if actual_doc_id is None and document_id: - actual_doc_id = document_id - - # Set document_id on the fact - processed_fact.document_id = actual_doc_id - - # Map chunk_id if this fact came from a chunk - if fact.chunk_index is not None: - # Look up chunk_id using (doc_id, chunk_index) - chunk_id = chunk_id_map_by_doc.get((actual_doc_id, fact.chunk_index)) - if chunk_id: - processed_fact.chunk_id = chunk_id - else: - # No chunks - still need to set document_id on facts - for fact, processed_fact in zip(extracted_facts, processed_facts): - original_doc_id = contents_dicts[fact.content_index].get("document_id") - # Map to actual document_id (handles None -> generated UUID mapping) - actual_doc_id = doc_id_mapping.get(original_doc_id, original_doc_id) - if actual_doc_id is None and document_id: - actual_doc_id = document_id - processed_fact.document_id = actual_doc_id - - non_duplicate_facts = processed_facts - - # Insert facts (document_id is now stored per-fact) - step_start = time.time() - unit_ids = await fact_storage.insert_facts_batch(conn, bank_id, non_duplicate_facts) - log_buffer.append(f"[5] Insert facts: {len(unit_ids)} units in {time.time() - step_start:.3f}s") - - # Process entities - step_start = time.time() - # Build map of content_index -> user entities for merging - user_entities_per_content = { - idx: content.entities for idx, content in enumerate(contents) if content.entities - } - entity_links = await entity_processing.process_entities_batch( - entity_resolver, - conn, - bank_id, - unit_ids, - non_duplicate_facts, - log_buffer, - user_entities_per_content=user_entities_per_content, - entity_labels=getattr(config, "entity_labels", None), - ) - log_buffer.append(f"[6] Process entities: {len(entity_links)} links in {time.time() - step_start:.3f}s") - - # Create temporal links - step_start = time.time() - temporal_link_count = await link_creation.create_temporal_links_batch(conn, bank_id, unit_ids) - log_buffer.append(f"[7] Temporal links: {temporal_link_count} links in {time.time() - step_start:.3f}s") - - # Create semantic links - step_start = time.time() - embeddings_for_links = [fact.embedding for fact in non_duplicate_facts] - semantic_link_count = await link_creation.create_semantic_links_batch( - conn, bank_id, unit_ids, embeddings_for_links - ) - log_buffer.append(f"[8] Semantic links: {semantic_link_count} links in {time.time() - step_start:.3f}s") - - # Insert entity links - step_start = time.time() - if entity_links: - await entity_processing.insert_entity_links_batch(conn, entity_links) - log_buffer.append( - f"[9] Entity links: {len(entity_links) if entity_links else 0} links in {time.time() - step_start:.3f}s" - ) - - # Create causal links - step_start = time.time() - causal_link_count = await link_creation.create_causal_links_batch(conn, unit_ids, non_duplicate_facts) - log_buffer.append(f"[10] Causal links: {causal_link_count} links in {time.time() - step_start:.3f}s") - - # Map results back to original content items - result_unit_ids = _map_results_to_contents(contents, extracted_facts, unit_ids) - - # Transactional outbox: queue any side-effect tasks (e.g. webhook deliveries) - # inside the same transaction so they are atomically committed with the retain data. - if outbox_callback: - await outbox_callback(conn) - - # Flush entity stats (mention_count / last_seen) now that the transaction - # has committed. Uses a fresh pool connection — no locks held. - await entity_resolver.flush_pending_stats() - - # Log final summary - total_time = time.time() - start_time - log_buffer.append(f"{'=' * 60}") - log_buffer.append(f"RETAIN_BATCH COMPLETE: {len(unit_ids)} units in {total_time:.3f}s") - if document_ids_added: - log_buffer.append(f"Documents: {', '.join(document_ids_added)}") - log_buffer.append(f"{'=' * 60}") - - logger.info("\n" + "\n".join(log_buffer) + "\n") - - await retry_with_backoff(_run_db_work) - return result_unit_ids, usage + return delta_contents, delta_chunk_map def _map_results_to_contents( diff --git a/hindsight-api-slim/hindsight_api/metrics.py b/hindsight-api-slim/hindsight_api/metrics.py index 996f07c6..1e29b7ac 100644 --- a/hindsight-api-slim/hindsight_api/metrics.py +++ b/hindsight-api-slim/hindsight_api/metrics.py @@ -11,14 +11,11 @@ This module provides metrics for: - Database connection pool metrics """ +import importlib import logging import os -import types -try: - import resource -except ImportError: - resource: types.ModuleType | None = None # Windows doesn't have resource module +_resource_mod = importlib.import_module("resource") if importlib.util.find_spec("resource") else None import threading import time from contextlib import contextmanager @@ -460,13 +457,13 @@ class MetricsCollector(MetricsCollectorBase): def _setup_process_metrics(self): """Set up observable gauges for process metrics.""" - if resource is None: + if _resource_mod is None: return # Skip process metrics on Windows def get_cpu_times(_options): """Get process CPU times.""" try: - rusage = resource.getrusage(resource.RUSAGE_SELF) + rusage = _resource_mod.getrusage(_resource_mod.RUSAGE_SELF) yield metrics.Observation(rusage.ru_utime, {"type": "user"}) yield metrics.Observation(rusage.ru_stime, {"type": "system"}) except Exception: @@ -475,7 +472,7 @@ class MetricsCollector(MetricsCollectorBase): def get_memory_usage(_options): """Get process memory usage in bytes.""" try: - rusage = resource.getrusage(resource.RUSAGE_SELF) + rusage = _resource_mod.getrusage(_resource_mod.RUSAGE_SELF) # ru_maxrss is in kilobytes on Linux, bytes on macOS max_rss = rusage.ru_maxrss if os.uname().sysname == "Linux": @@ -493,7 +490,7 @@ class MetricsCollector(MetricsCollectorBase): yield metrics.Observation(count) else: # Fallback: use resource limits - soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE) + soft, hard = _resource_mod.getrlimit(_resource_mod.RLIMIT_NOFILE) yield metrics.Observation(soft, {"limit": "soft"}) except Exception: pass diff --git a/hindsight-api-slim/tests/test_delta_retain.py b/hindsight-api-slim/tests/test_delta_retain.py new file mode 100644 index 00000000..675caf40 --- /dev/null +++ b/hindsight-api-slim/tests/test_delta_retain.py @@ -0,0 +1,842 @@ +""" +Tests for delta retain — upsert optimization that only re-processes changed chunks. +""" + +import logging +from datetime import datetime, timezone + +import pytest + +from hindsight_api import RequestContext +from hindsight_api.engine.memory_engine import Budget + +logger = logging.getLogger(__name__) + + +def _ts(): + return datetime.now(timezone.utc).timestamp() + + +# ============================================================ +# Core Delta Retain Tests +# ============================================================ + + +@pytest.mark.asyncio +async def test_delta_retain_unchanged_content_skips_llm(memory, request_context): + """ + When upserting a document with identical content, no new facts should be + extracted (LLM is not called for unchanged chunks). The existing facts + should be preserved. + """ + bank_id = f"test_delta_unchanged_{_ts()}" + document_id = "conversation-001" + + try: + content = "Alice works at Google. Bob works at Microsoft." + + # First retain — full processing + v1_units = await memory.retain_async( + bank_id=bank_id, + content=content, + context="team info", + document_id=document_id, + request_context=request_context, + ) + assert len(v1_units) > 0, "v1 should create facts" + + # Get v1 document state + doc_v1 = await memory.get_document(document_id, bank_id, request_context=request_context) + v1_unit_count = doc_v1["memory_unit_count"] + + # Second retain — same content, should use delta path (no new facts) + v2_units = await memory.retain_async( + bank_id=bank_id, + content=content, + context="team info", + document_id=document_id, + request_context=request_context, + ) + + # No new units should be returned (nothing changed) + assert v2_units == [], "Delta retain with unchanged content should return empty unit list" + + # Existing facts should still be there + doc_v2 = await memory.get_document(document_id, bank_id, request_context=request_context) + assert doc_v2["memory_unit_count"] == v1_unit_count, "Existing facts should be preserved" + + # Verify recall still works + result = await memory.recall_async( + bank_id=bank_id, + query="Where does Alice work?", + budget=Budget.MID, + max_tokens=1000, + request_context=request_context, + ) + assert len(result.results) > 0, "Should still recall facts after delta retain" + + finally: + await memory.delete_bank(bank_id, request_context=request_context) + + +@pytest.mark.asyncio +async def test_delta_retain_appended_content(memory, request_context): + """ + When a conversation grows (new content appended), only new chunks should + be processed. Facts from unchanged chunks should be preserved. + """ + bank_id = f"test_delta_append_{_ts()}" + document_id = "growing-conversation" + + try: + # First version — short content (single chunk) + v1_content = "Alice is a software engineer at Google. She works on search infrastructure." + + v1_units = await memory.retain_async( + bank_id=bank_id, + content=v1_content, + context="profile", + document_id=document_id, + request_context=request_context, + ) + assert len(v1_units) > 0 + + # Get v1 facts via recall + v1_recall = await memory.recall_async( + bank_id=bank_id, + query="What does Alice do?", + budget=Budget.MID, + max_tokens=2000, + request_context=request_context, + ) + v1_fact_texts = {r.text for r in v1_recall.results} + + # Second version — original content + new content appended + # This should preserve facts from the first chunk and add new ones + v2_content = v1_content + "\n\nBob joined Google as a product manager in 2024. He previously worked at Meta on AR/VR products." + + v2_units = await memory.retain_async( + bank_id=bank_id, + content=v2_content, + context="profile", + document_id=document_id, + request_context=request_context, + ) + + # Should have facts about Bob from the new content + v2_recall = await memory.recall_async( + bank_id=bank_id, + query="What does Bob do?", + budget=Budget.MID, + max_tokens=2000, + request_context=request_context, + ) + bob_facts = [r for r in v2_recall.results if "bob" in r.text.lower()] + assert len(bob_facts) > 0, "Should have facts about Bob from appended content" + + # Should still have facts about Alice from original content + alice_recall = await memory.recall_async( + bank_id=bank_id, + query="What does Alice do?", + budget=Budget.MID, + max_tokens=2000, + request_context=request_context, + ) + assert len(alice_recall.results) > 0, "Should still have Alice facts from original content" + + finally: + await memory.delete_bank(bank_id, request_context=request_context) + + +@pytest.mark.asyncio +async def test_delta_retain_modified_chunk(memory, request_context): + """ + When content in the middle changes, that chunk should be re-processed + while other chunks are preserved. + """ + bank_id = f"test_delta_modified_{_ts()}" + document_id = "changing-doc" + + try: + # v1: Alice works at Google + v1_content = "Alice works at Google as a senior engineer." + v1_units = await memory.retain_async( + bank_id=bank_id, + content=v1_content, + context="team", + document_id=document_id, + request_context=request_context, + ) + assert len(v1_units) > 0 + + # v2: Alice works at Microsoft (changed) + v2_content = "Alice works at Microsoft as a principal engineer." + v2_units = await memory.retain_async( + bank_id=bank_id, + content=v2_content, + context="team", + document_id=document_id, + request_context=request_context, + ) + + # New facts should reflect the updated content + result = await memory.recall_async( + bank_id=bank_id, + query="Where does Alice work?", + budget=Budget.MID, + max_tokens=2000, + request_context=request_context, + ) + all_texts = " ".join(r.text.lower() for r in result.results) + assert "microsoft" in all_texts, f"Should have updated fact about Microsoft, got: {all_texts}" + + finally: + await memory.delete_bank(bank_id, request_context=request_context) + + +# ============================================================ +# Entity & Link Tests +# ============================================================ + + +@pytest.mark.asyncio +async def test_delta_retain_entities_preserved_for_unchanged_chunks(memory, request_context): + """ + Entities linked to unchanged chunks should be preserved after delta retain. + """ + bank_id = f"test_delta_entities_{_ts()}" + document_id = "entity-doc" + + try: + v1_content = "Alice works at Google. She is a senior engineer in the Cloud division." + v1_units = await memory.retain_async( + bank_id=bank_id, + content=v1_content, + context="team", + document_id=document_id, + request_context=request_context, + ) + assert len(v1_units) > 0 + + # Check entities exist + pool = await memory._get_pool() + async with pool.acquire() as conn: + v1_entities = await conn.fetch( + "SELECT canonical_name FROM entities WHERE bank_id = $1", + bank_id, + ) + v1_entity_names = {e["canonical_name"].lower() for e in v1_entities} + assert len(v1_entity_names) > 0, "Should have entities after v1 retain" + + # Upsert with same content — entities should persist + await memory.retain_async( + bank_id=bank_id, + content=v1_content, + context="team", + document_id=document_id, + request_context=request_context, + ) + + async with pool.acquire() as conn: + v2_entities = await conn.fetch( + "SELECT canonical_name FROM entities WHERE bank_id = $1", + bank_id, + ) + v2_entity_names = {e["canonical_name"].lower() for e in v2_entities} + + # All v1 entities should still exist + assert v1_entity_names.issubset(v2_entity_names), ( + f"v1 entities {v1_entity_names} should be preserved, got {v2_entity_names}" + ) + + finally: + await memory.delete_bank(bank_id, request_context=request_context) + + +@pytest.mark.asyncio +async def test_delta_retain_new_entities_created_for_new_chunks(memory, request_context): + """ + New entities should be created for newly added chunks during delta retain. + """ + bank_id = f"test_delta_new_entities_{_ts()}" + document_id = "entity-growth-doc" + + try: + v1_content = "Alice works at Google." + await memory.retain_async( + bank_id=bank_id, + content=v1_content, + context="team", + document_id=document_id, + request_context=request_context, + ) + + pool = await memory._get_pool() + async with pool.acquire() as conn: + v1_entities = await conn.fetch( + "SELECT canonical_name FROM entities WHERE bank_id = $1", + bank_id, + ) + v1_entity_names = {e["canonical_name"].lower() for e in v1_entities} + + # Append content mentioning new entities + v2_content = v1_content + "\n\nBob joined Facebook. He works with Charlie on the Reality Labs project." + await memory.retain_async( + bank_id=bank_id, + content=v2_content, + context="team", + document_id=document_id, + request_context=request_context, + ) + + async with pool.acquire() as conn: + v2_entities = await conn.fetch( + "SELECT canonical_name FROM entities WHERE bank_id = $1", + bank_id, + ) + v2_entity_names = {e["canonical_name"].lower() for e in v2_entities} + + # Should have more entities after adding content with new people/orgs + assert len(v2_entity_names) > len(v1_entity_names), ( + f"Should have more entities after append: v1={v1_entity_names}, v2={v2_entity_names}" + ) + + finally: + await memory.delete_bank(bank_id, request_context=request_context) + + +@pytest.mark.asyncio +async def test_delta_retain_links_preserved_for_unchanged_chunks(memory, request_context): + """ + Memory links (temporal, semantic, entity) for unchanged chunks should be preserved. + """ + bank_id = f"test_delta_links_{_ts()}" + document_id = "links-doc" + + try: + content = "Alice is a senior engineer at Google Cloud. She mentors junior engineers and reviews their code." + v1_units = await memory.retain_async( + bank_id=bank_id, + content=content, + context="team", + document_id=document_id, + request_context=request_context, + ) + assert len(v1_units) > 0 + + # Count links after v1 + pool = await memory._get_pool() + async with pool.acquire() as conn: + v1_link_count = await conn.fetchval( + """SELECT COUNT(*) FROM memory_links ml + JOIN memory_units mu ON ml.from_unit_id = mu.id + WHERE mu.bank_id = $1 AND mu.document_id = $2""", + bank_id, + document_id, + ) + + # Upsert with same content + await memory.retain_async( + bank_id=bank_id, + content=content, + context="team", + document_id=document_id, + request_context=request_context, + ) + + # Links should be preserved + async with pool.acquire() as conn: + v2_link_count = await conn.fetchval( + """SELECT COUNT(*) FROM memory_links ml + JOIN memory_units mu ON ml.from_unit_id = mu.id + WHERE mu.bank_id = $1 AND mu.document_id = $2""", + bank_id, + document_id, + ) + + assert v2_link_count == v1_link_count, ( + f"Links should be preserved: v1={v1_link_count}, v2={v2_link_count}" + ) + + finally: + await memory.delete_bank(bank_id, request_context=request_context) + + +# ============================================================ +# Document Metadata & Tags Tests +# ============================================================ + + +@pytest.mark.asyncio +async def test_delta_retain_document_metadata_updated(memory, request_context): + """ + Document metadata (retain_params, tags) should be updated even when + chunk content hasn't changed. + """ + bank_id = f"test_delta_meta_{_ts()}" + document_id = "metadata-doc" + + try: + content = "Alice works at Google." + + # v1 with initial tags + await memory.retain_async( + bank_id=bank_id, + content=content, + context="initial context", + document_id=document_id, + request_context=request_context, + ) + + doc_v1 = await memory.get_document(document_id, bank_id, request_context=request_context) + assert doc_v1 is not None + + # v2 with updated context (same content — triggers delta path) + await memory.retain_async( + bank_id=bank_id, + content=content, + context="updated context", + document_id=document_id, + request_context=request_context, + ) + + doc_v2 = await memory.get_document(document_id, bank_id, request_context=request_context) + assert doc_v2 is not None + assert doc_v2["updated_at"] >= doc_v1["updated_at"], "Document should have updated timestamp" + + finally: + await memory.delete_bank(bank_id, request_context=request_context) + + +@pytest.mark.asyncio +async def test_delta_retain_tags_propagated_to_existing_units(memory, request_context): + """ + When tags change during an upsert with unchanged content, the new tags + should be propagated to all existing memory units. + """ + bank_id = f"test_delta_tags_{_ts()}" + document_id = "tags-doc" + + try: + content = "Alice works at Google." + + # v1 with tag "team-a" + await memory.retain_batch_async( + bank_id=bank_id, + contents=[{ + "content": content, + "document_id": document_id, + "tags": ["team-a"], + }], + request_context=request_context, + ) + + pool = await memory._get_pool() + async with pool.acquire() as conn: + v1_tags = await conn.fetch( + "SELECT tags FROM memory_units WHERE bank_id = $1 AND document_id = $2", + bank_id, + document_id, + ) + assert all("team-a" in row["tags"] for row in v1_tags), "v1 units should have team-a tag" + + # v2 with same content but different tags + await memory.retain_batch_async( + bank_id=bank_id, + contents=[{ + "content": content, + "document_id": document_id, + "tags": ["team-b", "important"], + }], + request_context=request_context, + ) + + async with pool.acquire() as conn: + v2_tags = await conn.fetch( + "SELECT tags FROM memory_units WHERE bank_id = $1 AND document_id = $2", + bank_id, + document_id, + ) + for row in v2_tags: + assert "team-b" in row["tags"], f"v2 units should have team-b tag, got {row['tags']}" + assert "important" in row["tags"], f"v2 units should have important tag, got {row['tags']}" + + finally: + await memory.delete_bank(bank_id, request_context=request_context) + + +# ============================================================ +# Chunk Management Tests +# ============================================================ + + +@pytest.mark.asyncio +async def test_delta_retain_removed_chunks_delete_facts(memory, request_context): + """ + When content is shortened (chunks removed), facts from the removed + chunks should be deleted. + """ + bank_id = f"test_delta_removed_{_ts()}" + document_id = "shrinking-doc" + + try: + # v1: longer content with facts about Alice and Bob + v1_content = ( + "Alice is a senior engineer at Google Cloud. " + "She leads the infrastructure team and has been there for 5 years.\n\n" + "Bob is a product manager at Facebook Reality Labs. " + "He previously worked at Amazon on Alexa voice products." + ) + + v1_units = await memory.retain_async( + bank_id=bank_id, + content=v1_content, + context="profiles", + document_id=document_id, + request_context=request_context, + ) + assert len(v1_units) > 0 + + doc_v1 = await memory.get_document(document_id, bank_id, request_context=request_context) + v1_count = doc_v1["memory_unit_count"] + + # v2: Completely different content — all chunks change + v2_content = "Charlie works at Netflix as a data scientist." + v2_units = await memory.retain_async( + bank_id=bank_id, + content=v2_content, + context="profiles", + document_id=document_id, + request_context=request_context, + ) + + doc_v2 = await memory.get_document(document_id, bank_id, request_context=request_context) + assert doc_v2 is not None + + # Should have facts about Charlie + result = await memory.recall_async( + bank_id=bank_id, + query="Who works at Netflix?", + budget=Budget.MID, + max_tokens=2000, + request_context=request_context, + ) + all_texts = " ".join(r.text.lower() for r in result.results) + assert "charlie" in all_texts or "netflix" in all_texts, ( + f"Should have facts about Charlie/Netflix after replacing content, got: {all_texts}" + ) + + finally: + await memory.delete_bank(bank_id, request_context=request_context) + + +@pytest.mark.asyncio +async def test_delta_retain_chunks_have_content_hash(memory, request_context): + """ + After retain, chunks should have content_hash populated. + """ + bank_id = f"test_delta_hash_{_ts()}" + document_id = "hash-doc" + + try: + content = "Alice works at Google as a software engineer." + await memory.retain_async( + bank_id=bank_id, + content=content, + document_id=document_id, + request_context=request_context, + ) + + pool = await memory._get_pool() + async with pool.acquire() as conn: + chunks = await conn.fetch( + "SELECT chunk_id, content_hash FROM chunks WHERE document_id = $1 AND bank_id = $2", + document_id, + bank_id, + ) + + assert len(chunks) > 0, "Should have stored chunks" + for chunk in chunks: + assert chunk["content_hash"] is not None, f"Chunk {chunk['chunk_id']} should have content_hash" + assert len(chunk["content_hash"]) == 64, "content_hash should be SHA256 hex (64 chars)" + + finally: + await memory.delete_bank(bank_id, request_context=request_context) + + +# ============================================================ +# Backward Compatibility Tests +# ============================================================ + + +@pytest.mark.asyncio +async def test_retain_without_document_id_still_works(memory, request_context): + """ + Retain without document_id should still work normally (no delta path). + """ + bank_id = f"test_no_docid_{_ts()}" + + try: + units = await memory.retain_async( + bank_id=bank_id, + content="Alice works at Google.", + context="test", + request_context=request_context, + ) + assert len(units) > 0, "Should create facts without document_id" + + result = await memory.recall_async( + bank_id=bank_id, + query="Where does Alice work?", + budget=Budget.MID, + max_tokens=1000, + request_context=request_context, + ) + assert len(result.results) > 0 + + finally: + await memory.delete_bank(bank_id, request_context=request_context) + + +@pytest.mark.asyncio +async def test_delta_retain_first_retain_full_path(memory, request_context): + """ + First retain of a new document should use the full path (no delta possible). + """ + bank_id = f"test_first_retain_{_ts()}" + document_id = "new-doc" + + try: + units = await memory.retain_async( + bank_id=bank_id, + content="Alice works at Google.", + context="test", + document_id=document_id, + request_context=request_context, + ) + assert len(units) > 0, "First retain should create facts via full path" + + doc = await memory.get_document(document_id, bank_id, request_context=request_context) + assert doc is not None + assert doc["memory_unit_count"] > 0 + + finally: + await memory.delete_bank(bank_id, request_context=request_context) + + +# ============================================================ +# Edge Cases +# ============================================================ + + +@pytest.mark.asyncio +async def test_delta_retain_empty_to_content(memory, request_context): + """ + Going from gibberish (zero facts) to real content should work. + """ + bank_id = f"test_delta_empty_{_ts()}" + document_id = "empty-to-content" + + try: + # v1: content that probably produces zero facts + await memory.retain_async( + bank_id=bank_id, + content="!!!###$$$%%%", + document_id=document_id, + request_context=request_context, + ) + + doc_v1 = await memory.get_document(document_id, bank_id, request_context=request_context) + assert doc_v1 is not None + + # v2: real content + v2_units = await memory.retain_async( + bank_id=bank_id, + content="Alice works at Google as a senior engineer.", + document_id=document_id, + request_context=request_context, + ) + + doc_v2 = await memory.get_document(document_id, bank_id, request_context=request_context) + assert doc_v2 is not None + assert doc_v2["memory_unit_count"] > 0 or len(v2_units) > 0, "Should have facts after updating with real content" + + finally: + await memory.delete_bank(bank_id, request_context=request_context) + + +@pytest.mark.asyncio +async def test_delta_retain_multiple_upserts(memory, request_context): + """ + Multiple sequential upserts should work correctly, with delta optimization + kicking in after the first retain. + """ + bank_id = f"test_delta_multi_{_ts()}" + document_id = "multi-upsert" + + try: + # v1: initial + v1_content = "Alice works at Google." + await memory.retain_async( + bank_id=bank_id, + content=v1_content, + document_id=document_id, + request_context=request_context, + ) + + # v2: same content (delta: no changes) + await memory.retain_async( + bank_id=bank_id, + content=v1_content, + document_id=document_id, + request_context=request_context, + ) + + # v3: append + v3_content = v1_content + "\n\nBob works at Microsoft." + await memory.retain_async( + bank_id=bank_id, + content=v3_content, + document_id=document_id, + request_context=request_context, + ) + + # v4: same as v3 (delta: no changes again) + await memory.retain_async( + bank_id=bank_id, + content=v3_content, + document_id=document_id, + request_context=request_context, + ) + + # Final check: should have facts about both Alice and Bob + result = await memory.recall_async( + bank_id=bank_id, + query="Who works where?", + budget=Budget.MID, + max_tokens=2000, + request_context=request_context, + ) + all_texts = " ".join(r.text.lower() for r in result.results) + assert "alice" in all_texts or "google" in all_texts, f"Should have Alice/Google facts, got: {all_texts}" + + doc = await memory.get_document(document_id, bank_id, request_context=request_context) + assert doc is not None + assert doc["memory_unit_count"] > 0 + + finally: + await memory.delete_bank(bank_id, request_context=request_context) + + +@pytest.mark.asyncio +async def test_delta_retain_with_user_entities(memory, request_context): + """ + User-provided entities should work correctly with delta retain. + """ + bank_id = f"test_delta_user_entities_{_ts()}" + document_id = "user-entity-doc" + + try: + content = "The project is going well." + + # v1 with user entities + await memory.retain_batch_async( + bank_id=bank_id, + contents=[{ + "content": content, + "document_id": document_id, + "entities": [{"text": "Project Alpha", "type": "PROJECT"}], + }], + request_context=request_context, + ) + + pool = await memory._get_pool() + async with pool.acquire() as conn: + v1_entities = await conn.fetch( + "SELECT canonical_name FROM entities WHERE bank_id = $1", + bank_id, + ) + v1_names = {e["canonical_name"].lower() for e in v1_entities} + + # v2 with additional entity, same content + # Note: same content = delta path (no re-extraction) + # The user entities for NEW chunks only get processed + v2_content = content + "\n\nThe timeline is on track for Q2 delivery." + await memory.retain_batch_async( + bank_id=bank_id, + contents=[{ + "content": v2_content, + "document_id": document_id, + "entities": [ + {"text": "Project Alpha", "type": "PROJECT"}, + {"text": "Q2 Deadline", "type": "MILESTONE"}, + ], + }], + request_context=request_context, + ) + + # Should have entities from both v1 and v2 + async with pool.acquire() as conn: + v2_entities = await conn.fetch( + "SELECT canonical_name FROM entities WHERE bank_id = $1", + bank_id, + ) + v2_names = {e["canonical_name"].lower() for e in v2_entities} + + # v1 entities should be preserved + assert v1_names.issubset(v2_names), f"v1 entities should be preserved: {v1_names} not in {v2_names}" + + finally: + await memory.delete_bank(bank_id, request_context=request_context) + + +@pytest.mark.asyncio +async def test_delta_retain_recall_with_chunks(memory, request_context): + """ + After delta retain, recall with include_chunks should return correct chunk data. + """ + bank_id = f"test_delta_recall_chunks_{_ts()}" + document_id = "recall-chunks-doc" + + try: + content = "Alice is a senior engineer at Google Cloud. She designs distributed systems." + await memory.retain_async( + bank_id=bank_id, + content=content, + context="profile", + document_id=document_id, + request_context=request_context, + ) + + # Upsert with same content (delta: no changes) + await memory.retain_async( + bank_id=bank_id, + content=content, + context="profile", + document_id=document_id, + request_context=request_context, + ) + + # Recall with chunks + result = await memory.recall_async( + bank_id=bank_id, + query="What does Alice do?", + budget=Budget.MID, + max_tokens=2000, + include_chunks=True, + max_chunk_tokens=8192, + request_context=request_context, + ) + + assert len(result.results) > 0, "Should recall facts" + + # Facts with chunk_ids should have corresponding chunks + facts_with_chunks = [r for r in result.results if r.chunk_id] + if facts_with_chunks and result.chunks: + for fact in facts_with_chunks: + assert fact.chunk_id in result.chunks, ( + f"Chunk {fact.chunk_id} should be in returned chunks" + ) + + finally: + await memory.delete_bank(bank_id, request_context=request_context) diff --git a/skills/hindsight-docs/references/changelog/integrations/claude-code.md b/skills/hindsight-docs/references/changelog/integrations/claude-code.md index 060d1446..9bef5aec 100644 --- a/skills/hindsight-docs/references/changelog/integrations/claude-code.md +++ b/skills/hindsight-docs/references/changelog/integrations/claude-code.md @@ -7,3 +7,27 @@ import PageHero from '@site/src/components/PageHero'; [← Claude Code integration](../../sdks/integrations/claude-code.md) + +## [0.3.0](https://github.com/vectorize-io/hindsight/tree/integrations/claude-code/v0.3.0) + +**Features** + +- Claude Code integration now retains tool calls as structured JSON for more accurate memory and retrieval. ([`8cb8b912`](https://github.com/vectorize-io/hindsight/commit/8cb8b912)) + +## [0.2.0](https://github.com/vectorize-io/hindsight/tree/integrations/claude-code/v0.2.0) + +**Features** + +- Added a Claude Code integration plugin for capturing and using Hindsight memory in Claude Code. ([`f4390bdc`](https://github.com/vectorize-io/hindsight/commit/f4390bdc)) +- Claude Code integration can retain full sessions with document upsert and configurable tagging. ([`2d31b67d`](https://github.com/vectorize-io/hindsight/commit/2d31b67d)) + +**Improvements** + +- Improved Claude Code plugin installation and configuration experience. ([`35b2cbb6`](https://github.com/vectorize-io/hindsight/commit/35b2cbb6)) +- Integrations no longer rely on hardcoded default models, allowing model selection to be fully configured. ([`58e68f3e`](https://github.com/vectorize-io/hindsight/commit/58e68f3e)) +- Claude Code now starts the Hindsight background daemon automatically at session start for smoother operation. ([`26944e25`](https://github.com/vectorize-io/hindsight/commit/26944e25)) + +**Bug Fixes** + +- Added a supported setup command to register hooks reliably, fixing hook registration issues. ([`22ca6a8d`](https://github.com/vectorize-io/hindsight/commit/22ca6a8d)) +- Fixed Claude Code integration compatibility on Windows. ([`a94a90ea`](https://github.com/vectorize-io/hindsight/commit/a94a90ea))