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

hindsight_api/engine/retain/types.py

@@ -6,8 +6,7 @@ from content input to fact storage.
 """
 
 from dataclasses import dataclass, field
-from typing import List, Optional, Dict, Any
-from datetime import datetime
+from datetime import UTC, datetime
 from uuid import UUID
 
 
@@ -18,16 +17,18 @@ class RetainContent:
 
     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)
+    event_date: datetime | None = 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)
+            from datetime import datetime
+
+            self.event_date = datetime.now(UTC)
 
 
 @dataclass
@@ -37,6 +38,7 @@ class ChunkMetadata:
 
     Used to track which facts were extracted from which chunks.
     """
+
     chunk_text: str
     fact_count: int
     content_index: int  # Index of the source content
@@ -50,9 +52,10 @@ class EntityRef:
 
     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
+    canonical_name: str | None = None  # Resolved canonical name
+    entity_id: UUID | None = None  # Resolved entity ID
 
 
 @dataclass
@@ -62,6 +65,7 @@ class CausalRelation:
 
     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
@@ -74,20 +78,21 @@ class ExtractedFact:
 
     This is the raw output from fact extraction before processing.
     """
+
     fact_text: str
     fact_type: str  # "world", "experience", "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)
+    entities: list[str] = field(default_factory=list)
+    occurred_start: datetime | None = None
+    occurred_end: datetime | None = None
+    where: str | None = 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)
+    mentioned_at: datetime | None = None
+    metadata: dict[str, str] = field(default_factory=dict)
 
 
 @dataclass
@@ -97,37 +102,38 @@ class ProcessedFact:
 
     Includes resolved entities, embeddings, and all necessary fields.
     """
+
     # Core fact data
     fact_text: str
     fact_type: str
-    embedding: List[float]
+    embedding: list[float]
 
     # Temporal data
-    occurred_start: Optional[datetime]
-    occurred_end: Optional[datetime]
+    occurred_start: datetime | None
+    occurred_end: datetime | None
     mentioned_at: datetime
 
     # Context and metadata
     context: str
-    metadata: Dict[str, str]
+    metadata: dict[str, str]
 
     # Location data
-    where: Optional[str] = None
+    where: str | None = None
 
     # Entities
-    entities: List[EntityRef] = field(default_factory=list)
+    entities: list[EntityRef] = field(default_factory=list)
 
     # Causal relations
-    causal_relations: List[CausalRelation] = field(default_factory=list)
+    causal_relations: list[CausalRelation] = field(default_factory=list)
 
     # Chunk reference
-    chunk_id: Optional[str] = None
+    chunk_id: str | None = None
 
     # Document reference (denormalized for query performance)
-    document_id: Optional[str] = None
+    document_id: str | None = None
 
     # DB fields (set after insertion)
-    unit_id: Optional[UUID] = None
+    unit_id: UUID | None = None
 
     @property
     def is_duplicate(self) -> bool:
@@ -136,10 +142,8 @@ class ProcessedFact:
 
     @staticmethod
     def from_extracted_fact(
-        extracted_fact: 'ExtractedFact',
-        embedding: List[float],
-        chunk_id: Optional[str] = None
-    ) -> 'ProcessedFact':
+        extracted_fact: "ExtractedFact", embedding: list[float], chunk_id: str | None = None
+    ) -> "ProcessedFact":
         """
         Create ProcessedFact from ExtractedFact.
 
@@ -151,12 +155,12 @@ class ProcessedFact:
         Returns:
             ProcessedFact ready for storage
         """
-        from datetime import datetime, timezone
+        from datetime import datetime
 
         # 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)
+        mentioned_at = extracted_fact.mentioned_at or datetime.now(UTC)
 
         # Convert entity strings to EntityRef objects
         entities = [EntityRef(name=name) for name in extracted_fact.entities]
@@ -172,7 +176,7 @@ class ProcessedFact:
             metadata=extracted_fact.metadata,
             entities=entities,
             causal_relations=extracted_fact.causal_relations,
-            chunk_id=chunk_id
+            chunk_id=chunk_id,
         )
 
 
@@ -183,10 +187,11 @@ class EntityLink:
 
     Used for entity-based graph connections in the memory graph.
     """
+
     from_unit_id: UUID
     to_unit_id: UUID
     entity_id: UUID
-    link_type: str = 'entity'
+    link_type: str = "entity"
     weight: float = 1.0
 
 
@@ -197,24 +202,25 @@ class RetainBatch:
 
     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
+    contents: list[RetainContent]
+    document_id: str | None = None
+    fact_type_override: str | None = None
+    confidence_score: float | None = 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)
+    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)
+    unit_ids_by_content: list[list[str]] = field(default_factory=list)
 
-    def get_facts_for_content(self, content_index: int) -> List[ExtractedFact]:
+    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]:
+    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

@@ -7,15 +7,15 @@ Provides modular search architecture:
 - Reranking: Pluggable strategies (heuristic, cross-encoder)
 """
 
+from .graph_retrieval import BFSGraphRetriever, GraphRetriever
+from .mpfp_retrieval import MPFPGraphRetriever
+from .reranking import CrossEncoderReranker
 from .retrieval import (
-    retrieve_parallel,
+    ParallelRetrievalResult,
     get_default_graph_retriever,
+    retrieve_parallel,
     set_default_graph_retriever,
-    ParallelRetrievalResult,
 )
-from .graph_retrieval import GraphRetriever, BFSGraphRetriever
-from .mpfp_retrieval import MPFPGraphRetriever
-from .reranking import CrossEncoderReranker
 
 __all__ = [
     "retrieve_parallel",

hindsight_api/engine/search/fusion.py

@@ -2,15 +2,12 @@
 Helper functions for hybrid search (semantic + BM25 + graph).
 """
 
-from typing import List, Dict, Any, Tuple
-import asyncio
-from .types import RetrievalResult, MergedCandidate
+from typing import Any
 
+from .types import MergedCandidate, RetrievalResult
 
-def reciprocal_rank_fusion(
-    result_lists: List[List[RetrievalResult]],
-    k: int = 60
-) -> List[MergedCandidate]:
+
+def reciprocal_rank_fusion(result_lists: list[list[RetrievalResult]], k: int = 60) -> list[MergedCandidate]:
     """
     Merge multiple ranked result lists using Reciprocal Rank Fusion.
 
@@ -73,20 +70,14 @@ def reciprocal_rank_fusion(
         sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True), start=1
     ):
         merged_candidate = MergedCandidate(
-            retrieval=all_retrievals[doc_id],
-            rrf_score=rrf_score,
-            rrf_rank=rrf_rank,
-            source_ranks=source_ranks[doc_id]
+            retrieval=all_retrievals[doc_id], rrf_score=rrf_score, rrf_rank=rrf_rank, source_ranks=source_ranks[doc_id]
         )
         merged_results.append(merged_candidate)
 
     return merged_results
 
 
-def normalize_scores_on_deltas(
-    results: List[Dict[str, Any]],
-    score_keys: List[str]
-) -> List[Dict[str, Any]]:
+def normalize_scores_on_deltas(results: list[dict[str, Any]], score_keys: list[str]) -> list[dict[str, Any]]:
     """
     Normalize scores based on deltas (min-max normalization within result set).
 

hindsight_api/engine/search/graph_retrieval.py

@@ -6,13 +6,11 @@ allowing different algorithms (BFS spreading activation, PPR, etc.) to be
 swapped without changing the rest of the recall pipeline.
 """
 
-from abc import ABC, abstractmethod
-from typing import List, Optional
-from datetime import datetime
 import logging
+from abc import ABC, abstractmethod
 
-from .types import RetrievalResult
 from ..db_utils import acquire_with_retry
+from .types import RetrievalResult
 
 logger = logging.getLogger(__name__)
 
@@ -40,10 +38,10 @@ class GraphRetriever(ABC):
         bank_id: str,
         fact_type: str,
         budget: int,
-        query_text: Optional[str] = None,
-        semantic_seeds: Optional[List[RetrievalResult]] = None,
-        temporal_seeds: Optional[List[RetrievalResult]] = None,
-    ) -> List[RetrievalResult]:
+        query_text: str | None = None,
+        semantic_seeds: list[RetrievalResult] | None = None,
+        temporal_seeds: list[RetrievalResult] | None = None,
+    ) -> list[RetrievalResult]:
         """
         Retrieve relevant facts via graph traversal.
 
@@ -109,10 +107,10 @@ class BFSGraphRetriever(GraphRetriever):
         bank_id: str,
         fact_type: str,
         budget: int,
-        query_text: Optional[str] = None,
-        semantic_seeds: Optional[List[RetrievalResult]] = None,
-        temporal_seeds: Optional[List[RetrievalResult]] = None,
-    ) -> List[RetrievalResult]:
+        query_text: str | None = None,
+        semantic_seeds: list[RetrievalResult] | None = None,
+        temporal_seeds: list[RetrievalResult] | None = None,
+    ) -> list[RetrievalResult]:
         """
         Retrieve facts using BFS spreading activation.
 
@@ -127,9 +125,7 @@ class BFSGraphRetriever(GraphRetriever):
         for interface compatibility but not used.
         """
         async with acquire_with_retry(pool) as conn:
-            return await self._retrieve_with_conn(
-                conn, query_embedding_str, bank_id, fact_type, budget
-            )
+            return await self._retrieve_with_conn(conn, query_embedding_str, bank_id, fact_type, budget)
 
     async def _retrieve_with_conn(
         self,
@@ -138,7 +134,7 @@ class BFSGraphRetriever(GraphRetriever):
         bank_id: str,
         fact_type: str,
         budget: int,
-    ) -> List[RetrievalResult]:
+    ) -> list[RetrievalResult]:
         """Internal implementation with connection."""
 
         # Step 1: Find entry points
@@ -155,8 +151,11 @@ class BFSGraphRetriever(GraphRetriever):
             ORDER BY embedding <=> $1::vector
             LIMIT $5
             """,
-            query_embedding_str, bank_id, fact_type,
-            self.entry_point_threshold, self.entry_point_limit
+            query_embedding_str,
+            bank_id,
+            fact_type,
+            self.entry_point_threshold,
+            self.entry_point_limit,
         )
 
         if not entry_points:
@@ -165,10 +164,7 @@ class BFSGraphRetriever(GraphRetriever):
         # Step 2: BFS spreading activation
         visited = set()
         results = []
-        queue = [
-            (RetrievalResult.from_db_row(dict(r)), r["similarity"])
-            for r in entry_points
-        ]
+        queue = [(RetrievalResult.from_db_row(dict(r)), r["similarity"]) for r in entry_points]
         budget_remaining = budget
 
         while queue and budget_remaining > 0:
@@ -205,7 +201,10 @@ class BFSGraphRetriever(GraphRetriever):
                     ORDER BY ml.weight DESC
                     LIMIT $4
                     """,
-                    batch_nodes, self.min_activation, fact_type, max_neighbors
+                    batch_nodes,
+                    self.min_activation,
+                    fact_type,
+                    max_neighbors,
                 )
 
                 for n in neighbors: