hindsight-api 0.0.13__py3-none-any.whl

hindsight_api/engine/retain/chunk_storage.py

@@ -0,0 +1,82 @@
+"""
+Chunk storage for retain pipeline.
+
+Handles storage of document chunks in the database.
+"""
+import logging
+from typing import List, Dict, Optional
+
+from .types import ChunkMetadata
+
+logger = logging.getLogger(__name__)
+
+
+async def store_chunks_batch(
+    conn,
+    bank_id: str,
+    document_id: str,
+    chunks: List[ChunkMetadata]
+) -> Dict[int, str]:
+    """
+    Store document chunks in the database.
+
+    Args:
+        conn: Database connection
+        bank_id: Bank identifier
+        document_id: Document identifier
+        chunks: List of ChunkMetadata objects
+
+    Returns:
+        Dictionary mapping global chunk index to chunk_id
+    """
+    if not chunks:
+        return {}
+
+    # Prepare chunk data for batch insert
+    chunk_ids = []
+    chunk_texts = []
+    chunk_indices = []
+    chunk_id_map = {}
+
+    for chunk in chunks:
+        chunk_id = f"{bank_id}_{document_id}_{chunk.chunk_index}"
+        chunk_ids.append(chunk_id)
+        chunk_texts.append(chunk.chunk_text)
+        chunk_indices.append(chunk.chunk_index)
+        chunk_id_map[chunk.chunk_index] = chunk_id
+
+    # Batch insert all chunks
+    await conn.execute(
+        """
+        INSERT INTO chunks (chunk_id, document_id, bank_id, chunk_text, chunk_index)
+        SELECT * FROM unnest($1::text[], $2::text[], $3::text[], $4::text[], $5::integer[])
+        """,
+        chunk_ids,
+        [document_id] * len(chunk_texts),
+        [bank_id] * len(chunk_texts),
+        chunk_texts,
+        chunk_indices
+    )
+
+    return chunk_id_map
+
+
+def map_facts_to_chunks(
+    facts_chunk_indices: List[int],
+    chunk_id_map: Dict[int, str]
+) -> List[Optional[str]]:
+    """
+    Map fact chunk indices to chunk IDs.
+
+    Args:
+        facts_chunk_indices: List of chunk indices for each fact
+        chunk_id_map: Dictionary mapping chunk index to chunk_id
+
+    Returns:
+        List of chunk_ids (same length as facts_chunk_indices)
+    """
+    chunk_ids = []
+    for chunk_idx in facts_chunk_indices:
+        chunk_id = chunk_id_map.get(chunk_idx)
+        chunk_ids.append(chunk_id)
+    return chunk_ids

hindsight_api/engine/retain/deduplication.py

@@ -0,0 +1,104 @@
+"""
+Deduplication logic for retain pipeline.
+
+Checks for duplicate facts using semantic similarity and temporal proximity.
+"""
+import logging
+from datetime import datetime
+from typing import List
+from collections import defaultdict
+
+from .types import ProcessedFact
+
+logger = logging.getLogger(__name__)
+
+
+async def check_duplicates_batch(
+    conn,
+    bank_id: str,
+    facts: List[ProcessedFact],
+    duplicate_checker_fn
+) -> List[bool]:
+    """
+    Check which facts are duplicates using batched time-window queries.
+
+    Groups facts by 12-hour time buckets to efficiently check for duplicates
+    within a 24-hour window.
+
+    Args:
+        conn: Database connection
+        bank_id: Bank identifier
+        facts: List of ProcessedFact objects to check
+        duplicate_checker_fn: Async function(conn, bank_id, texts, embeddings, date, time_window_hours)
+                              that returns List[bool] indicating duplicates
+
+    Returns:
+        List of boolean flags (same length as facts) indicating if each fact is a duplicate
+    """
+    if not facts:
+        return []
+
+    # Group facts by event_date (rounded to 12-hour buckets) for efficient batching
+    time_buckets = defaultdict(list)
+    for idx, fact in enumerate(facts):
+        # Use occurred_start if available, otherwise use mentioned_at
+        # For deduplication purposes, we need a time reference
+        fact_date = fact.occurred_start if fact.occurred_start is not None else fact.mentioned_at
+
+        # Defensive: if both are None (shouldn't happen), use now()
+        if fact_date is None:
+            from datetime import datetime, timezone
+            fact_date = datetime.now(timezone.utc)
+
+        # Round to 12-hour bucket to group similar times
+        bucket_key = fact_date.replace(
+            hour=(fact_date.hour // 12) * 12,
+            minute=0,
+            second=0,
+            microsecond=0
+        )
+        time_buckets[bucket_key].append((idx, fact))
+
+    # Process each bucket in batch
+    all_is_duplicate = [False] * len(facts)
+
+    for bucket_date, bucket_items in time_buckets.items():
+        indices = [item[0] for item in bucket_items]
+        texts = [item[1].fact_text for item in bucket_items]
+        embeddings = [item[1].embedding for item in bucket_items]
+
+        # Check duplicates for this time bucket
+        dup_flags = await duplicate_checker_fn(
+            conn,
+            bank_id,
+            texts,
+            embeddings,
+            bucket_date,
+            time_window_hours=24
+        )
+
+        # Map results back to original indices
+        for idx, is_dup in zip(indices, dup_flags):
+            all_is_duplicate[idx] = is_dup
+
+    return all_is_duplicate
+
+
+def filter_duplicates(
+    facts: List[ProcessedFact],
+    is_duplicate_flags: List[bool]
+) -> List[ProcessedFact]:
+    """
+    Filter out duplicate facts based on duplicate flags.
+
+    Args:
+        facts: List of ProcessedFact objects
+        is_duplicate_flags: Boolean flags indicating which facts are duplicates
+
+    Returns:
+        List of non-duplicate facts
+    """
+    if len(facts) != len(is_duplicate_flags):
+        raise ValueError(f"Mismatch between facts ({len(facts)}) and flags ({len(is_duplicate_flags)})")
+
+    return [fact for fact, is_dup in zip(facts, is_duplicate_flags) if not is_dup]

hindsight_api/engine/retain/embedding_processing.py

@@ -0,0 +1,62 @@
+"""
+Embedding processing for retain pipeline.
+
+Handles augmenting fact texts with temporal information and generating embeddings.
+"""
+import logging
+from typing import List
+from datetime import datetime
+
+from . import embedding_utils
+from .types import ExtractedFact
+
+logger = logging.getLogger(__name__)
+
+
+def augment_texts_with_dates(facts: List[ExtractedFact], format_date_fn) -> List[str]:
+    """
+    Augment fact texts with readable dates for better temporal matching.
+
+    This allows queries like "camping in June" to match facts that happened in June.
+
+    Args:
+        facts: List of ExtractedFact objects
+        format_date_fn: Function to format datetime to readable string
+
+    Returns:
+        List of augmented text strings (same length as facts)
+    """
+    augmented_texts = []
+    for fact in facts:
+        # Use occurred_start as the representative date
+        fact_date = fact.occurred_start or fact.mentioned_at
+        readable_date = format_date_fn(fact_date)
+        # Augment text with date for embedding (but store original text in DB)
+        augmented_text = f"{fact.fact_text} (happened in {readable_date})"
+        augmented_texts.append(augmented_text)
+    return augmented_texts
+
+
+async def generate_embeddings_batch(
+    embeddings_model,
+    texts: List[str]
+) -> List[List[float]]:
+    """
+    Generate embeddings for a batch of texts.
+
+    Args:
+        embeddings_model: Embeddings model instance
+        texts: List of text strings to embed
+
+    Returns:
+        List of embedding vectors (same length as texts)
+    """
+    if not texts:
+        return []
+
+    embeddings = await embedding_utils.generate_embeddings_batch(
+        embeddings_model,
+        texts
+    )
+
+    return embeddings

hindsight_api/engine/retain/embedding_utils.py

@@ -0,0 +1,54 @@
+"""
+Embedding generation utilities for memory units.
+"""
+
+import asyncio
+import logging
+from typing import List
+
+logger = logging.getLogger(__name__)
+
+
+def generate_embedding(embeddings_backend, text: str) -> List[float]:
+    """
+    Generate embedding for text using the provided embeddings backend.
+
+    Args:
+        embeddings_backend: Embeddings instance to use for encoding
+        text: Text to embed
+
+    Returns:
+        Embedding vector (dimension depends on embeddings backend)
+    """
+    try:
+        embeddings = embeddings_backend.encode([text])
+        return embeddings[0]
+    except Exception as e:
+        raise Exception(f"Failed to generate embedding: {str(e)}")
+
+
+async def generate_embeddings_batch(embeddings_backend, texts: List[str]) -> List[List[float]]:
+    """
+    Generate embeddings for multiple texts using the provided embeddings backend.
+
+    Runs the embedding generation in a thread pool to avoid blocking the event loop
+    for CPU-bound operations.
+
+    Args:
+        embeddings_backend: Embeddings instance to use for encoding
+        texts: List of texts to embed
+
+    Returns:
+        List of embeddings in same order as input texts
+    """
+    try:
+        # Run embeddings in thread pool to avoid blocking event loop
+        loop = asyncio.get_event_loop()
+        embeddings = await loop.run_in_executor(
+            None,  # Use default thread pool
+            embeddings_backend.encode,
+            texts
+        )
+        return embeddings
+    except Exception as e:
+        raise Exception(f"Failed to generate batch embeddings: {str(e)}")

hindsight_api/engine/retain/entity_processing.py

@@ -0,0 +1,90 @@
+"""
+Entity processing for retain pipeline.
+
+Handles entity extraction, resolution, and link creation for stored facts.
+"""
+import logging
+from typing import List, Tuple, Dict, Any
+from uuid import UUID
+
+from .types import ProcessedFact, EntityRef
+from . import link_utils
+
+logger = logging.getLogger(__name__)
+
+
+async def process_entities_batch(
+    entity_resolver,
+    conn,
+    bank_id: str,
+    unit_ids: List[str],
+    facts: List[ProcessedFact],
+    log_buffer: List[str] = None
+) -> List[Tuple[str, str, float]]:
+    """
+    Process entities for all facts and create entity links.
+
+    This function:
+    1. Extracts entity mentions from fact texts
+    2. Resolves entity names to canonical entities
+    3. Creates entity records in the database
+    4. Returns entity links ready for insertion
+
+    Args:
+        entity_resolver: EntityResolver instance for entity resolution
+        conn: Database connection
+        bank_id: Bank identifier
+        unit_ids: List of unit IDs (same length as facts)
+        facts: List of ProcessedFact objects
+        log_buffer: Optional buffer for detailed logging
+
+    Returns:
+        List of entity link tuples: (unit_id, entity_id, confidence)
+    """
+    if not unit_ids or not facts:
+        return []
+
+    if len(unit_ids) != len(facts):
+        raise ValueError(f"Mismatch between unit_ids ({len(unit_ids)}) and facts ({len(facts)})")
+
+    # Extract data for link_utils function
+    fact_texts = [fact.fact_text for fact in facts]
+    # Use occurred_start if available, otherwise use mentioned_at for entity timestamps
+    fact_dates = [fact.occurred_start if fact.occurred_start is not None else fact.mentioned_at for fact in facts]
+    # Convert EntityRef objects to dict format expected by link_utils
+    entities_per_fact = [
+        [{'text': entity.name, 'type': 'CONCEPT'} for entity in (fact.entities or [])]
+        for fact in facts
+    ]
+
+    # Use existing link_utils function for entity processing
+    entity_links = await link_utils.extract_entities_batch_optimized(
+        entity_resolver,
+        conn,
+        bank_id,
+        unit_ids,
+        fact_texts,
+        "",  # context (not used in current implementation)
+        fact_dates,
+        entities_per_fact,
+        log_buffer  # Pass log_buffer for detailed logging
+    )
+
+    return entity_links
+
+
+async def insert_entity_links_batch(
+    conn,
+    entity_links: List[Tuple[str, str, float]]
+) -> None:
+    """
+    Insert entity links in batch.
+
+    Args:
+        conn: Database connection
+        entity_links: List of (unit_id, entity_id, confidence) tuples
+    """
+    if not entity_links:
+        return
+
+    await link_utils.insert_entity_links_batch(conn, entity_links)