hindsight-api 0.1.5__py3-none-any.whl → 0.1.7__py3-none-any.whl

hindsight_api/engine/retain/fact_storage.py

@@ -3,10 +3,9 @@ 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
+import logging
 
 from .types import ProcessedFact
 
@@ -14,11 +13,8 @@ logger = logging.getLogger(__name__)
 
 
 async def insert_facts_batch(
-    conn,
-    bank_id: str,
-    facts: List[ProcessedFact],
-    document_id: Optional[str] = None
-) -> List[str]:
+    conn, bank_id: str, facts: list[ProcessedFact], document_id: str | None = None
+) -> list[str]:
     """
     Insert facts into the database in batch.
 
@@ -62,7 +58,7 @@ async def insert_facts_batch(
         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)
+        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)
@@ -93,10 +89,10 @@ async def insert_facts_batch(
         access_counts,
         metadata_jsons,
         chunk_ids,
-        document_ids
+        document_ids,
     )
 
-    unit_ids = [str(row['id']) for row in results]
+    unit_ids = [str(row["id"]) for row in results]
     return unit_ids
 
 
@@ -119,17 +115,12 @@ async def ensure_bank_exists(conn, bank_id: str) -> None:
         """,
         bank_id,
         '{"skepticism": 3, "literalism": 3, "empathy": 3}',
-        ""
+        "",
     )
 
 
 async def handle_document_tracking(
-    conn,
-    bank_id: str,
-    document_id: str,
-    combined_content: str,
-    is_first_batch: bool,
-    retain_params: Optional[dict] = None
+    conn, bank_id: str, document_id: str, combined_content: str, is_first_batch: bool, retain_params: dict | None = None
 ) -> None:
     """
     Handle document tracking in the database.
@@ -150,10 +141,7 @@ async def handle_document_tracking(
     # 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
-        )
+        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(
@@ -172,5 +160,5 @@ async def handle_document_tracking(
         combined_content,
         content_hash,
         json.dumps({}),  # Empty metadata dict
-        json.dumps(retain_params) if retain_params else None
+        json.dumps(retain_params) if retain_params else None,
     )

hindsight_api/engine/retain/link_creation.py

@@ -3,20 +3,16 @@ 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
+from .types import ProcessedFact
 
 logger = logging.getLogger(__name__)
 
 
-async def create_temporal_links_batch(
-    conn,
-    bank_id: str,
-    unit_ids: List[str]
-) -> int:
+async def create_temporal_links_batch(conn, bank_id: str, unit_ids: list[str]) -> int:
     """
     Create temporal links between facts.
 
@@ -33,20 +29,10 @@ async def create_temporal_links_batch(
     if not unit_ids:
         return 0
 
-    return await link_utils.create_temporal_links_batch_per_fact(
-        conn,
-        bank_id,
-        unit_ids,
-        log_buffer=[]
-    )
+    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]]
-) -> int:
+async def create_semantic_links_batch(conn, bank_id: str, unit_ids: list[str], embeddings: list[list[float]]) -> int:
     """
     Create semantic links between facts.
 
@@ -67,20 +53,10 @@ async def create_semantic_links_batch(
     if len(unit_ids) != len(embeddings):
         raise ValueError(f"Mismatch between unit_ids ({len(unit_ids)}) and embeddings ({len(embeddings)})")
 
-    return await link_utils.create_semantic_links_batch(
-        conn,
-        bank_id,
-        unit_ids,
-        embeddings,
-        log_buffer=[]
-    )
+    return 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:
+async def create_causal_links_batch(conn, unit_ids: list[str], facts: list[ProcessedFact]) -> int:
     """
     Create causal links between facts.
 
@@ -108,9 +84,9 @@ async def create_causal_links_batch(
             # Convert CausalRelation objects to dicts
             relations_dicts = [
                 {
-                    'relation_type': rel.relation_type,
-                    'target_fact_index': rel.target_fact_index,
-                    'strength': rel.strength
+                    "relation_type": rel.relation_type,
+                    "target_fact_index": rel.target_fact_index,
+                    "strength": rel.strength,
                 }
                 for rel in fact.causal_relations
             ]
@@ -118,10 +94,6 @@ async def create_causal_links_batch(
         else:
             causal_relations_per_fact.append([])
 
-    link_count = await link_utils.create_causal_links_batch(
-        conn,
-        unit_ids,
-        causal_relations_per_fact
-    )
+    link_count = await link_utils.create_causal_links_batch(conn, unit_ids, causal_relations_per_fact)
 
     return link_count