hindsight-api 0.0.13__py3-none-any.whl

hindsight_api/engine/retain/orchestrator.py

@@ -0,0 +1,405 @@
+"""
+Main orchestrator for the retain pipeline.
+
+Coordinates all retain pipeline modules to store memories efficiently.
+"""
+import logging
+import time
+import uuid
+from datetime import datetime, timezone
+from typing import List, Dict, Any, Optional
+
+from . import bank_utils
+from ..db_utils import acquire_with_retry
+
+
+def utcnow():
+    """Get current UTC time."""
+    return datetime.now(timezone.utc)
+
+from .types import RetainContent, ExtractedFact, ProcessedFact
+from . import (
+    fact_extraction,
+    embedding_processing,
+    deduplication,
+    chunk_storage,
+    fact_storage,
+    entity_processing,
+    link_creation
+)
+
+logger = logging.getLogger(__name__)
+
+
+async def retain_batch(
+    pool,
+    embeddings_model,
+    llm_config,
+    entity_resolver,
+    task_backend,
+    format_date_fn,
+    duplicate_checker_fn,
+    regenerate_observations_fn,
+    bank_id: str,
+    contents_dicts: List[Dict[str, Any]],
+    document_id: Optional[str] = None,
+    is_first_batch: bool = True,
+    fact_type_override: Optional[str] = None,
+    confidence_score: Optional[float] = None,
+) -> List[List[str]]:
+    """
+    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
+        task_backend: Task backend for background jobs
+        format_date_fn: Function to format datetime to readable string
+        duplicate_checker_fn: Function to check for duplicate facts
+        regenerate_observations_fn: Async function to regenerate observations for entities
+        bank_id: Bank identifier
+        contents_dicts: List of content dictionaries
+        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
+
+    Returns:
+        List of unit ID lists (one list per content item)
+    """
+    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}")
+    log_buffer.append(f"Batch size: {len(contents_dicts)} content items, {total_chars:,} chars")
+    log_buffer.append(f"{'='*60}")
+
+    # Get bank profile
+    profile = await bank_utils.get_bank_profile(pool, bank_id)
+    agent_name = profile["name"]
+
+    # Convert dicts to RetainContent objects
+    contents = []
+    for item in contents_dicts:
+        content = RetainContent(
+            content=item["content"],
+            context=item.get("context", ""),
+            event_date=item.get("event_date") or utcnow(),
+            metadata=item.get("metadata", {})
+        )
+        contents.append(content)
+
+    # Step 1: Extract facts from all contents
+    step_start = time.time()
+    extract_opinions = (fact_type_override == 'opinion')
+
+    extracted_facts, chunks = await fact_extraction.extract_facts_from_contents(
+        contents,
+        llm_config,
+        agent_name,
+        extract_opinions
+    )
+    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")
+
+    if not extracted_facts:
+        return [[] for _ in contents]
+
+    # 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)
+    ]
+
+    # Track document IDs for logging
+    document_ids_added = []
+
+    # 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
+    async with acquire_with_retry(pool) as conn:
+        async with conn.transaction():
+            # Ensure bank exists
+            await fact_storage.ensure_bank_exists(conn, bank_id)
+
+            # 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 = {}
+                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
+                )
+                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])
+
+                            # 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
+                            )
+                            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
+
+            # Deduplication
+            step_start = time.time()
+            is_duplicate_flags = await deduplication.check_duplicates_batch(
+                conn, bank_id, processed_facts, duplicate_checker_fn
+            )
+            log_buffer.append(f"[4] Deduplication: {sum(is_duplicate_flags)} duplicates in {time.time() - step_start:.3f}s")
+
+            # Filter out duplicates
+            non_duplicate_facts = deduplication.filter_duplicates(processed_facts, is_duplicate_flags)
+
+            if not non_duplicate_facts:
+                return [[] for _ in contents]
+
+            # 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()
+            entity_links = await entity_processing.process_entities_batch(
+                entity_resolver, conn, bank_id, unit_ids, non_duplicate_facts, log_buffer
+            )
+            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()
+            await link_creation.create_temporal_links_batch(conn, bank_id, unit_ids)
+            log_buffer.append(f"[7] Temporal links: {time.time() - step_start:.3f}s")
+
+            # Create semantic links
+            step_start = time.time()
+            embeddings_for_links = [fact.embedding for fact in non_duplicate_facts]
+            await link_creation.create_semantic_links_batch(conn, bank_id, unit_ids, embeddings_for_links)
+            log_buffer.append(f"[8] Semantic links: {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: {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, is_duplicate_flags, unit_ids
+            )
+
+            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")
+
+        # Trigger background tasks AFTER transaction commits
+        await _trigger_background_tasks(
+            task_backend,
+            regenerate_observations_fn,
+            bank_id,
+            unit_ids,
+            non_duplicate_facts,
+            entity_links
+        )
+
+        return result_unit_ids
+
+
+def _map_results_to_contents(
+    contents: List[RetainContent],
+    extracted_facts: List[ExtractedFact],
+    is_duplicate_flags: List[bool],
+    unit_ids: List[str]
+) -> List[List[str]]:
+    """
+    Map created unit IDs back to original content items.
+
+    Accounts for duplicates when mapping back.
+    """
+    result_unit_ids = []
+    filtered_idx = 0
+
+    # Group facts by content_index
+    facts_by_content = {i: [] for i in range(len(contents))}
+    for i, fact in enumerate(extracted_facts):
+        facts_by_content[fact.content_index].append(i)
+
+    for content_index in range(len(contents)):
+        content_unit_ids = []
+        for fact_idx in facts_by_content[content_index]:
+            if not is_duplicate_flags[fact_idx]:
+                content_unit_ids.append(unit_ids[filtered_idx])
+                filtered_idx += 1
+        result_unit_ids.append(content_unit_ids)
+
+    return result_unit_ids
+
+
+async def _trigger_background_tasks(
+    task_backend,
+    regenerate_observations_fn,
+    bank_id: str,
+    unit_ids: List[str],
+    facts: List[ProcessedFact],
+    entity_links: List
+) -> None:
+    """Trigger opinion reinforcement and observation regeneration (sync)."""
+    # Trigger opinion reinforcement if there are entities
+    fact_entities = [[e.name for e in fact.entities] for fact in facts]
+    if any(fact_entities):
+        await task_backend.submit_task({
+            'type': 'reinforce_opinion',
+            'bank_id': bank_id,
+            'created_unit_ids': unit_ids,
+            'unit_texts': [fact.fact_text for fact in facts],
+            'unit_entities': fact_entities
+        })
+
+    # Regenerate observations synchronously for top entities
+    TOP_N_ENTITIES = 5
+    MIN_FACTS_THRESHOLD = 5
+
+    if entity_links and regenerate_observations_fn:
+        unique_entity_ids = set()
+        for link in entity_links:
+            # links are tuples: (unit_id, entity_id, confidence)
+            if len(link) >= 2 and link[1]:
+                unique_entity_ids.add(str(link[1]))
+
+        if unique_entity_ids:
+            # Run observation regeneration synchronously
+            await regenerate_observations_fn(
+                bank_id=bank_id,
+                entity_ids=list(unique_entity_ids)[:TOP_N_ENTITIES],
+                min_facts=MIN_FACTS_THRESHOLD
+            )

hindsight_api/engine/retain/types.py

@@ -0,0 +1,206 @@
+"""
+Type definitions for the retain pipeline.
+
+These dataclasses provide type safety throughout the retain operation,
+from content input to fact storage.
+"""
+
+from dataclasses import dataclass, field
+from typing import List, Optional, Dict, Any
+from datetime import datetime
+from uuid import UUID
+
+
+@dataclass
+class RetainContent:
+    """
+    Input content item to be retained as memories.
+
+    Represents a single piece of content to extract facts from.
+    """
+    content: str
+    context: str = ""
+    event_date: Optional[datetime] = None
+    metadata: Dict[str, str] = field(default_factory=dict)
+
+    def __post_init__(self):
+        """Ensure event_date is set."""
+        if self.event_date is None:
+            from datetime import datetime, timezone
+            self.event_date = datetime.now(timezone.utc)
+
+
+@dataclass
+class ChunkMetadata:
+    """
+    Metadata about a text chunk.
+
+    Used to track which facts were extracted from which chunks.
+    """
+    chunk_text: str
+    fact_count: int
+    content_index: int  # Index of the source content
+    chunk_index: int  # Global chunk index across all contents
+
+
+@dataclass
+class EntityRef:
+    """
+    Reference to an entity mentioned in a fact.
+
+    Entities are extracted by the LLM during fact extraction.
+    """
+    name: str
+    canonical_name: Optional[str] = None  # Resolved canonical name
+    entity_id: Optional[UUID] = None  # Resolved entity ID
+
+
+@dataclass
+class CausalRelation:
+    """
+    Causal relationship between facts.
+
+    Represents how one fact causes, enables, or prevents another.
+    """
+    relation_type: str  # "causes", "enables", "prevents", "caused_by"
+    target_fact_index: int  # Index of the target fact in the batch
+    strength: float = 1.0  # Strength of the causal relationship
+
+
+@dataclass
+class ExtractedFact:
+    """
+    Fact extracted from content by the LLM.
+
+    This is the raw output from fact extraction before processing.
+    """
+    fact_text: str
+    fact_type: str  # "world", "bank", "opinion", "observation"
+    entities: List[str] = field(default_factory=list)
+    occurred_start: Optional[datetime] = None
+    occurred_end: Optional[datetime] = None
+    where: Optional[str] = None  # WHERE the fact occurred or is about
+    causal_relations: List[CausalRelation] = field(default_factory=list)
+
+    # Context from the content item
+    content_index: int = 0  # Which content this fact came from
+    chunk_index: int = 0  # Which chunk this fact came from
+    context: str = ""
+    mentioned_at: Optional[datetime] = None
+    metadata: Dict[str, str] = field(default_factory=dict)
+
+
+@dataclass
+class ProcessedFact:
+    """
+    Fact after processing and ready for storage.
+
+    Includes resolved entities, embeddings, and all necessary fields.
+    """
+    # Core fact data
+    fact_text: str
+    fact_type: str
+    embedding: List[float]
+
+    # Temporal data
+    occurred_start: Optional[datetime]
+    occurred_end: Optional[datetime]
+    mentioned_at: datetime
+
+    # Context and metadata
+    context: str
+    metadata: Dict[str, str]
+
+    # Location data
+    where: Optional[str] = None
+
+    # Entities
+    entities: List[EntityRef] = field(default_factory=list)
+
+    # Causal relations
+    causal_relations: List[CausalRelation] = field(default_factory=list)
+
+    # Chunk reference
+    chunk_id: Optional[str] = None
+
+    # Document reference (denormalized for query performance)
+    document_id: Optional[str] = None
+
+    # DB fields (set after insertion)
+    unit_id: Optional[UUID] = None
+
+    @property
+    def is_duplicate(self) -> bool:
+        """Check if this fact was marked as a duplicate."""
+        return self.unit_id is None
+
+    @staticmethod
+    def from_extracted_fact(
+        extracted_fact: 'ExtractedFact',
+        embedding: List[float],
+        chunk_id: Optional[str] = None
+    ) -> 'ProcessedFact':
+        """
+        Create ProcessedFact from ExtractedFact.
+
+        Args:
+            extracted_fact: Source ExtractedFact
+            embedding: Generated embedding vector
+            chunk_id: Optional chunk ID
+
+        Returns:
+            ProcessedFact ready for storage
+        """
+        from datetime import datetime, timezone
+
+        # Use occurred dates only if explicitly provided by LLM
+        occurred_start = extracted_fact.occurred_start
+        occurred_end = extracted_fact.occurred_end
+        mentioned_at = extracted_fact.mentioned_at or datetime.now(timezone.utc)
+
+        # Convert entity strings to EntityRef objects
+        entities = [EntityRef(name=name) for name in extracted_fact.entities]
+
+        return ProcessedFact(
+            fact_text=extracted_fact.fact_text,
+            fact_type=extracted_fact.fact_type,
+            embedding=embedding,
+            occurred_start=occurred_start,
+            occurred_end=occurred_end,
+            mentioned_at=mentioned_at,
+            context=extracted_fact.context,
+            metadata=extracted_fact.metadata,
+            entities=entities,
+            causal_relations=extracted_fact.causal_relations,
+            chunk_id=chunk_id
+        )
+
+
+@dataclass
+class RetainBatch:
+    """
+    A batch of content to retain.
+
+    Tracks all facts, chunks, and metadata for a batch operation.
+    """
+    bank_id: str
+    contents: List[RetainContent]
+    document_id: Optional[str] = None
+    fact_type_override: Optional[str] = None
+    confidence_score: Optional[float] = None
+
+    # Extracted data (populated during processing)
+    extracted_facts: List[ExtractedFact] = field(default_factory=list)
+    processed_facts: List[ProcessedFact] = field(default_factory=list)
+    chunks: List[ChunkMetadata] = field(default_factory=list)
+
+    # Results (populated after storage)
+    unit_ids_by_content: List[List[str]] = field(default_factory=list)
+
+    def get_facts_for_content(self, content_index: int) -> List[ExtractedFact]:
+        """Get all extracted facts for a specific content item."""
+        return [f for f in self.extracted_facts if f.content_index == content_index]
+
+    def get_chunks_for_content(self, content_index: int) -> List[ChunkMetadata]:
+        """Get all chunks for a specific content item."""
+        return [c for c in self.chunks if c.content_index == content_index]

hindsight_api/engine/search/__init__.py

@@ -0,0 +1,15 @@
+"""
+Search module for memory retrieval.
+
+Provides modular search architecture:
+- Retrieval: 4-way parallel (semantic + BM25 + graph + temporal)
+- Reranking: Pluggable strategies (heuristic, cross-encoder)
+"""
+
+from .retrieval import retrieve_parallel
+from .reranking import CrossEncoderReranker
+
+__all__ = [
+    "retrieve_parallel",
+    "CrossEncoderReranker",
+]