hindsight-api 0.0.13__py3-none-any.whl

hindsight_api/engine/retain/fact_storage.py

@@ -0,0 +1,176 @@
+"""
+Fact storage for retain pipeline.
+
+Handles insertion of facts into the database.
+"""
+import logging
+import json
+from typing import List, Optional
+from uuid import UUID
+
+from .types import ProcessedFact
+
+logger = logging.getLogger(__name__)
+
+
+async def insert_facts_batch(
+    conn,
+    bank_id: str,
+    facts: List[ProcessedFact],
+    document_id: Optional[str] = None
+) -> List[str]:
+    """
+    Insert facts into the database in batch.
+
+    Args:
+        conn: Database connection
+        bank_id: Bank identifier
+        facts: List of ProcessedFact objects to insert
+        document_id: Optional document ID to associate with facts
+
+    Returns:
+        List of unit IDs (UUIDs as strings) for the inserted facts
+    """
+    if not facts:
+        return []
+
+    # Prepare data for batch insert
+    fact_texts = []
+    embeddings = []
+    event_dates = []
+    occurred_starts = []
+    occurred_ends = []
+    mentioned_ats = []
+    contexts = []
+    fact_types = []
+    confidence_scores = []
+    access_counts = []
+    metadata_jsons = []
+    chunk_ids = []
+    document_ids = []
+
+    for fact in facts:
+        fact_texts.append(fact.fact_text)
+        # Convert embedding to string for asyncpg vector type
+        embeddings.append(str(fact.embedding))
+        # event_date: Use occurred_start if available, otherwise use mentioned_at
+        # This maintains backward compatibility while handling None occurred_start
+        event_dates.append(fact.occurred_start if fact.occurred_start is not None else fact.mentioned_at)
+        occurred_starts.append(fact.occurred_start)
+        occurred_ends.append(fact.occurred_end)
+        mentioned_ats.append(fact.mentioned_at)
+        contexts.append(fact.context)
+        fact_types.append(fact.fact_type)
+        # confidence_score is only for opinion facts
+        confidence_scores.append(1.0 if fact.fact_type == 'opinion' else None)
+        access_counts.append(0)  # Initial access count
+        metadata_jsons.append(json.dumps(fact.metadata))
+        chunk_ids.append(fact.chunk_id)
+        # Use per-fact document_id if available, otherwise fallback to batch-level document_id
+        document_ids.append(fact.document_id if fact.document_id else document_id)
+
+    # Batch insert all facts
+    results = await conn.fetch(
+        """
+        INSERT INTO memory_units (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
+                                 context, fact_type, confidence_score, access_count, metadata, chunk_id, document_id)
+        SELECT $1, * FROM unnest(
+            $2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
+            $8::text[], $9::text[], $10::float[], $11::int[], $12::jsonb[], $13::text[], $14::text[]
+        )
+        RETURNING id
+        """,
+        bank_id,
+        fact_texts,
+        embeddings,
+        event_dates,  # event_date: occurred_start if available, else mentioned_at
+        occurred_starts,
+        occurred_ends,
+        mentioned_ats,
+        contexts,
+        fact_types,
+        confidence_scores,
+        access_counts,
+        metadata_jsons,
+        chunk_ids,
+        document_ids
+    )
+
+    unit_ids = [str(row['id']) for row in results]
+    return unit_ids
+
+
+async def ensure_bank_exists(conn, bank_id: str) -> None:
+    """
+    Ensure bank exists in the database.
+
+    Creates bank with default values if it doesn't exist.
+
+    Args:
+        conn: Database connection
+        bank_id: Bank identifier
+    """
+    await conn.execute(
+        """
+        INSERT INTO banks (bank_id, personality, background)
+        VALUES ($1, $2::jsonb, $3)
+        ON CONFLICT (bank_id) DO UPDATE
+        SET updated_at = NOW()
+        """,
+        bank_id,
+        '{"openness": 0.5, "conscientiousness": 0.5, "extraversion": 0.5, "agreeableness": 0.5, "neuroticism": 0.5, "bias_strength": 0.5}',
+        ""
+    )
+
+
+async def handle_document_tracking(
+    conn,
+    bank_id: str,
+    document_id: str,
+    combined_content: str,
+    is_first_batch: bool,
+    retain_params: Optional[dict] = None
+) -> None:
+    """
+    Handle document tracking in the database.
+
+    Args:
+        conn: Database connection
+        bank_id: Bank identifier
+        document_id: Document identifier
+        combined_content: Combined content text from all content items
+        is_first_batch: Whether this is the first batch (for chunked operations)
+        retain_params: Optional parameters passed during retain (context, event_date, etc.)
+    """
+    import hashlib
+
+    # Calculate content hash
+    content_hash = hashlib.sha256(combined_content.encode()).hexdigest()
+
+    # Always delete old document first if it exists (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(
+            "DELETE FROM documents WHERE id = $1 AND bank_id = $2 RETURNING id",
+            document_id, bank_id
+        )
+
+    # Insert document (or update if exists from concurrent operations)
+    await conn.execute(
+        """
+        INSERT INTO documents (id, bank_id, original_text, content_hash, metadata, retain_params)
+        VALUES ($1, $2, $3, $4, $5, $6)
+        ON CONFLICT (id, bank_id) DO UPDATE
+        SET original_text = EXCLUDED.original_text,
+            content_hash = EXCLUDED.content_hash,
+            metadata = EXCLUDED.metadata,
+            retain_params = EXCLUDED.retain_params,
+            updated_at = NOW()
+        """,
+        document_id,
+        bank_id,
+        combined_content,
+        content_hash,
+        json.dumps({}),  # Empty metadata dict
+        json.dumps(retain_params) if retain_params else None
+    )

hindsight_api/engine/retain/link_creation.py

@@ -0,0 +1,121 @@
+"""
+Link creation for retain pipeline.
+
+Handles creation of temporal, semantic, and causal links between facts.
+"""
+import logging
+from typing import List
+
+from .types import ProcessedFact, CausalRelation
+from . import link_utils
+
+logger = logging.getLogger(__name__)
+
+
+async def create_temporal_links_batch(
+    conn,
+    bank_id: str,
+    unit_ids: List[str]
+) -> None:
+    """
+    Create temporal links between facts.
+
+    Links facts that occurred close in time to each other.
+
+    Args:
+        conn: Database connection
+        bank_id: Bank identifier
+        unit_ids: List of unit IDs to create links for
+    """
+    if not unit_ids:
+        return
+
+    await link_utils.create_temporal_links_batch_per_fact(
+        conn,
+        bank_id,
+        unit_ids,
+        log_buffer=[]
+    )
+
+
+async def create_semantic_links_batch(
+    conn,
+    bank_id: str,
+    unit_ids: List[str],
+    embeddings: List[List[float]]
+) -> None:
+    """
+    Create semantic links between facts.
+
+    Links facts that are semantically similar based on embeddings.
+
+    Args:
+        conn: Database connection
+        bank_id: Bank identifier
+        unit_ids: List of unit IDs to create links for
+        embeddings: List of embedding vectors (same length as unit_ids)
+    """
+    if not unit_ids or not embeddings:
+        return
+
+    if len(unit_ids) != len(embeddings):
+        raise ValueError(f"Mismatch between unit_ids ({len(unit_ids)}) and embeddings ({len(embeddings)})")
+
+    await link_utils.create_semantic_links_batch(
+        conn,
+        bank_id,
+        unit_ids,
+        embeddings,
+        log_buffer=[]
+    )
+
+
+async def create_causal_links_batch(
+    conn,
+    unit_ids: List[str],
+    facts: List[ProcessedFact]
+) -> int:
+    """
+    Create causal links between facts.
+
+    Links facts that have causal relationships (causes, enables, prevents).
+
+    Args:
+        conn: Database connection
+        unit_ids: List of unit IDs (same length as facts)
+        facts: List of ProcessedFact objects with causal_relations
+
+    Returns:
+        Number of causal links created
+    """
+    if not unit_ids or not facts:
+        return 0
+
+    if len(unit_ids) != len(facts):
+        raise ValueError(f"Mismatch between unit_ids ({len(unit_ids)}) and facts ({len(facts)})")
+
+    # Extract causal relations in the format expected by link_utils
+    # Format: List of lists, where each inner list is the causal relations for that fact
+    causal_relations_per_fact = []
+    for fact in facts:
+        if fact.causal_relations:
+            # Convert CausalRelation objects to dicts
+            relations_dicts = [
+                {
+                    'relation_type': rel.relation_type,
+                    'target_fact_index': rel.target_fact_index,
+                    'strength': rel.strength
+                }
+                for rel in fact.causal_relations
+            ]
+            causal_relations_per_fact.append(relations_dicts)
+        else:
+            causal_relations_per_fact.append([])
+
+    link_count = await link_utils.create_causal_links_batch(
+        conn,
+        unit_ids,
+        causal_relations_per_fact
+    )
+
+    return link_count