hindsight-api 0.2.1__py3-none-any.whl → 0.4.0__py3-none-any.whl

hindsight_api/engine/task_backend.py

@@ -1,31 +1,40 @@
 """
-Abstract task backend for running async tasks.
+Task backend for distributed task processing.
 
-This provides an abstraction that can be adapted to different execution models:
-- AsyncIO queue (default implementation)
-- Pub/Sub architectures (future)
-- Message brokers (future)
+This provides an abstraction for task storage and execution:
+- BrokerTaskBackend: Uses PostgreSQL as broker (production)
+- SyncTaskBackend: Executes tasks immediately (testing/embedded)
 """
 
-import asyncio
+import json
 import logging
 from abc import ABC, abstractmethod
 from collections.abc import Awaitable, Callable
-from typing import Any
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    import asyncpg
 
 logger = logging.getLogger(__name__)
 
 
+def fq_table(table: str, schema: str | None = None) -> str:
+    """Get fully-qualified table name with optional schema prefix."""
+    if schema:
+        return f'"{schema}".{table}'
+    return table
+
+
 class TaskBackend(ABC):
     """
     Abstract base class for task execution backends.
 
     Implementations must:
     1. Store/publish task events (as serializable dicts)
-    2. Execute tasks through a provided executor callback
+    2. Execute tasks through a provided executor callback (optional)
 
     The backend treats tasks as pure dictionaries that can be serialized
-    and sent over the network. The executor (typically MemoryEngine.execute_task)
+    and stored in the database. The executor (typically MemoryEngine.execute_task)
     receives the dict and routes it to the appropriate handler.
     """
 
@@ -46,7 +55,7 @@ class TaskBackend(ABC):
     @abstractmethod
     async def initialize(self):
         """
-        Initialize the backend (e.g., start workers, connect to broker).
+        Initialize the backend (e.g., connect to database).
         """
         pass
 
@@ -63,7 +72,7 @@ class TaskBackend(ABC):
     @abstractmethod
     async def shutdown(self):
         """
-        Shutdown the backend gracefully (e.g., stop workers, close connections).
+        Shutdown the backend gracefully.
         """
         pass
 
@@ -93,9 +102,8 @@ class SyncTaskBackend(TaskBackend):
     """
     Synchronous task backend that executes tasks immediately.
 
-    This is useful for embedded/CLI usage where we don't want background
-    workers that prevent clean exit. Tasks are executed inline rather than
-    being queued.
+    This is useful for tests and embedded/CLI usage where we don't want
+    background workers. Tasks are executed inline rather than being queued.
     """
 
     async def initialize(self):
@@ -121,130 +129,129 @@ class SyncTaskBackend(TaskBackend):
         logger.debug("SyncTaskBackend shutdown")
 
 
-class AsyncIOQueueBackend(TaskBackend):
+class BrokerTaskBackend(TaskBackend):
     """
-    Task backend implementation using asyncio queues.
+    Task backend using PostgreSQL as broker.
 
-    This is the default implementation that uses in-process asyncio queues
-    and a periodic consumer worker.
+    submit_task() stores task_payload in async_operations table.
+    Actual polling and execution is handled separately by WorkerPoller.
+
+    This backend is used by the API to store tasks. Workers poll
+    the database separately to claim and execute tasks.
     """
 
-    def __init__(self, batch_size: int = 100, batch_interval: float = 1.0):
+    def __init__(
+        self,
+        pool_getter: Callable[[], "asyncpg.Pool"],
+        schema: str | None = None,
+        schema_getter: Callable[[], str | None] | None = None,
+    ):
         """
-        Initialize AsyncIO queue backend.
+        Initialize the broker task backend.
 
         Args:
-            batch_size: Maximum number of tasks to process in one batch
-            batch_interval: Maximum time (seconds) to wait before processing batch
+            pool_getter: Callable that returns the asyncpg connection pool
+            schema: Database schema for multi-tenant support (optional, static)
+            schema_getter: Callable that returns current schema dynamically (optional).
+                          If set, takes precedence over static schema for submit_task.
         """
         super().__init__()
-        self._queue: asyncio.Queue | None = None
-        self._worker_task: asyncio.Task | None = None
-        self._shutdown_event: asyncio.Event | None = None
-        self._batch_size = batch_size
-        self._batch_interval = batch_interval
+        self._pool_getter = pool_getter
+        self._schema = schema
+        self._schema_getter = schema_getter
 
     async def initialize(self):
-        """Initialize the queue and start the worker."""
-        if self._initialized:
-            return
-
-        self._queue = asyncio.Queue()
-        self._shutdown_event = asyncio.Event()
-        self._worker_task = asyncio.create_task(self._worker())
+        """Initialize the backend."""
         self._initialized = True
-        logger.info("AsyncIOQueueBackend initialized")
+        logger.info("BrokerTaskBackend initialized")
 
     async def submit_task(self, task_dict: dict[str, Any]):
         """
-        Submit a task by putting it in the queue.
+        Store task payload in async_operations table.
+
+        The task_dict should contain an 'operation_id' if updating an existing
+        operation record, otherwise a new operation will be created.
 
         Args:
-            task_dict: Task dictionary to execute
+            task_dict: Task dictionary to store (must be JSON serializable)
         """
         if not self._initialized:
             await self.initialize()
 
-        await self._queue.put(task_dict)
+        pool = self._pool_getter()
+        operation_id = task_dict.get("operation_id")
         task_type = task_dict.get("type", "unknown")
-        task_id = task_dict.get("id")
+        bank_id = task_dict.get("bank_id")
+        payload_json = json.dumps(task_dict)
+
+        schema = self._schema_getter() if self._schema_getter else self._schema
+        table = fq_table("async_operations", schema)
+
+        if operation_id:
+            # Update existing operation with task payload
+            await pool.execute(
+                f"""
+                UPDATE {table}
+                SET task_payload = $1::jsonb, updated_at = now()
+                WHERE operation_id = $2
+                """,
+                payload_json,
+                operation_id,
+            )
+            logger.debug(f"Updated task payload for operation {operation_id}")
+        else:
+            # Insert new operation (for tasks without pre-created records)
+            # e.g., access_count_update tasks
+            import uuid
+
+            new_id = uuid.uuid4()
+            await pool.execute(
+                f"""
+                INSERT INTO {table} (operation_id, bank_id, operation_type, status, task_payload)
+                VALUES ($1, $2, $3, 'pending', $4::jsonb)
+                """,
+                new_id,
+                bank_id,
+                task_type,
+                payload_json,
+            )
+            logger.debug(f"Created new operation {new_id} for task type {task_type}")
 
-    async def wait_for_pending_tasks(self, timeout: float = 5.0):
+    async def shutdown(self):
+        """Shutdown the backend."""
+        self._initialized = False
+        logger.info("BrokerTaskBackend shutdown")
+
+    async def wait_for_pending_tasks(self, timeout: float = 120.0):
         """
-        Wait for all pending tasks in the queue to be processed.
+        Wait for pending tasks to be processed.
 
-        This is useful in tests to ensure background tasks complete before assertions.
+        In the broker model, this polls the database to check if tasks
+        for this process have been completed. This is useful in tests
+        when worker_enabled=True (API processes its own tasks).
 
         Args:
             timeout: Maximum time to wait in seconds
         """
-        if not self._initialized or self._queue is None:
-            return
+        import asyncio
+
+        pool = self._pool_getter()
+        schema = self._schema_getter() if self._schema_getter else self._schema
+        table = fq_table("async_operations", schema)
 
-        # Wait for queue to be empty and give worker time to process
         start_time = asyncio.get_event_loop().time()
         while asyncio.get_event_loop().time() - start_time < timeout:
-            if self._queue.empty():
-                # Queue is empty, give worker a bit more time to finish any in-flight task
-                await asyncio.sleep(0.3)
-                # Check again - if still empty, we're done
-                if self._queue.empty():
-                    return
-            else:
-                # Queue not empty, wait a bit
-                await asyncio.sleep(0.1)
+            # Check if there are any pending tasks with payloads
+            count = await pool.fetchval(
+                f"""
+                SELECT COUNT(*) FROM {table}
+                WHERE status = 'pending' AND task_payload IS NOT NULL
+                """
+            )
 
-    async def shutdown(self):
-        """Shutdown the worker and drain the queue."""
-        if not self._initialized:
-            return
-
-        logger.info("Shutting down AsyncIOQueueBackend...")
+            if count == 0:
+                return
 
-        # Signal shutdown
-        self._shutdown_event.set()
+            await asyncio.sleep(0.5)
 
-        # Cancel worker
-        if self._worker_task is not None:
-            self._worker_task.cancel()
-            try:
-                await self._worker_task
-            except asyncio.CancelledError:
-                pass  # Worker cancelled successfully
-
-        self._initialized = False
-        logger.info("AsyncIOQueueBackend shutdown complete")
-
-    async def _worker(self):
-        """
-        Background worker that processes tasks in batches.
-
-        Collects tasks for up to batch_interval seconds or batch_size items,
-        then processes them.
-        """
-        while not self._shutdown_event.is_set():
-            try:
-                # Collect tasks for batching
-                tasks = []
-                deadline = asyncio.get_event_loop().time() + self._batch_interval
-
-                while len(tasks) < self._batch_size and asyncio.get_event_loop().time() < deadline:
-                    try:
-                        remaining_time = max(0.1, deadline - asyncio.get_event_loop().time())
-                        task_dict = await asyncio.wait_for(self._queue.get(), timeout=remaining_time)
-                        tasks.append(task_dict)
-                    except TimeoutError:
-                        break
-
-                # Process batch
-                if tasks:
-                    # Execute tasks concurrently
-                    await asyncio.gather(
-                        *[self._execute_task(task_dict) for task_dict in tasks], return_exceptions=True
-                    )
-
-            except asyncio.CancelledError:
-                break
-            except Exception as e:
-                logger.error(f"Worker error: {e}")
-                await asyncio.sleep(1)  # Backoff on error
+        logger.warning(f"Timeout waiting for pending tasks after {timeout}s")

hindsight_api/engine/utils.py

@@ -49,7 +49,7 @@ async def extract_facts(
     if not text or not text.strip():
         return [], []
 
-    facts, chunks = await extract_facts_from_text(
+    facts, chunks, _ = await extract_facts_from_text(
         text,
         event_date,
         context=context,
@@ -65,154 +65,3 @@ async def extract_facts(
         return [], chunks
 
     return facts, chunks
-
-
-def cosine_similarity(vec1: list[float], vec2: list[float]) -> float:
-    """
-    Calculate cosine similarity between two vectors.
-
-    Args:
-        vec1: First vector
-        vec2: Second vector
-
-    Returns:
-        Similarity score between 0 and 1
-    """
-    if len(vec1) != len(vec2):
-        raise ValueError("Vectors must have same dimension")
-
-    dot_product = sum(a * b for a, b in zip(vec1, vec2))
-    magnitude1 = sum(a * a for a in vec1) ** 0.5
-    magnitude2 = sum(b * b for b in vec2) ** 0.5
-
-    if magnitude1 == 0 or magnitude2 == 0:
-        return 0.0
-
-    return dot_product / (magnitude1 * magnitude2)
-
-
-def calculate_recency_weight(days_since: float, half_life_days: float = 365.0) -> float:
-    """
-    Calculate recency weight using logarithmic decay.
-
-    This provides much better differentiation over long time periods compared to
-    exponential decay. Uses a log-based decay where the half-life parameter controls
-    when memories reach 50% weight.
-
-    Examples:
-        - Today (0 days): 1.0
-        - 1 year (365 days): ~0.5 (with default half_life=365)
-        - 2 years (730 days): ~0.33
-        - 5 years (1825 days): ~0.17
-        - 10 years (3650 days): ~0.09
-
-    This ensures that 2-year-old and 5-year-old memories have meaningfully
-    different weights, unlike exponential decay which makes them both ~0.
-
-    Args:
-        days_since: Number of days since the memory was created
-        half_life_days: Number of days for weight to reach 0.5 (default: 1 year)
-
-    Returns:
-        Weight between 0 and 1
-    """
-    import math
-
-    # Logarithmic decay: 1 / (1 + log(1 + days_since/half_life))
-    # This decays much slower than exponential, giving better long-term differentiation
-    normalized_age = days_since / half_life_days
-    return 1.0 / (1.0 + math.log1p(normalized_age))
-
-
-def calculate_frequency_weight(access_count: int, max_boost: float = 2.0) -> float:
-    """
-    Calculate frequency weight based on access count.
-
-    Frequently accessed memories are weighted higher.
-    Uses logarithmic scaling to avoid over-weighting.
-
-    Args:
-        access_count: Number of times the memory was accessed
-        max_boost: Maximum multiplier for frequently accessed memories
-
-    Returns:
-        Weight between 1.0 and max_boost
-    """
-    import math
-
-    if access_count <= 0:
-        return 1.0
-
-    # Logarithmic scaling: log(access_count + 1) / log(10)
-    # This gives: 0 accesses = 1.0, 9 accesses ~= 1.5, 99 accesses ~= 2.0
-    normalized = math.log(access_count + 1) / math.log(10)
-    return 1.0 + min(normalized, max_boost - 1.0)
-
-
-def calculate_temporal_anchor(occurred_start: datetime, occurred_end: datetime) -> datetime:
-    """
-    Calculate a single temporal anchor point from a temporal range.
-
-    Used for spreading activation - we need a single representative date
-    to calculate temporal proximity between facts. This simplifies the
-    range-to-range distance problem.
-
-    Strategy: Use midpoint of the range for balanced representation.
-
-    Args:
-        occurred_start: Start of temporal range
-        occurred_end: End of temporal range
-
-    Returns:
-        Single datetime representing the temporal anchor (midpoint)
-
-    Examples:
-        - Point event (July 14): start=July 14, end=July 14 → anchor=July 14
-        - Month range (February): start=Feb 1, end=Feb 28 → anchor=Feb 14
-        - Year range (2023): start=Jan 1, end=Dec 31 → anchor=July 1
-    """
-    # Calculate midpoint
-    time_delta = occurred_end - occurred_start
-    midpoint = occurred_start + (time_delta / 2)
-    return midpoint
-
-
-def calculate_temporal_proximity(anchor_a: datetime, anchor_b: datetime, half_life_days: float = 30.0) -> float:
-    """
-    Calculate temporal proximity between two temporal anchors.
-
-    Used for spreading activation to determine how "close" two facts are
-    in time. Uses logarithmic decay so that temporal similarity doesn't
-    drop off too quickly.
-
-    Args:
-        anchor_a: Temporal anchor of first fact
-        anchor_b: Temporal anchor of second fact
-        half_life_days: Number of days for proximity to reach 0.5
-                       (default: 30 days = 1 month)
-
-    Returns:
-        Proximity score in [0, 1] where:
-        - 1.0 = same day
-        - 0.5 = ~half_life days apart
-        - 0.0 = very distant in time
-
-    Examples:
-        - Same day: 1.0
-        - 1 week apart (half_life=30): ~0.7
-        - 1 month apart (half_life=30): ~0.5
-        - 1 year apart (half_life=30): ~0.2
-    """
-    import math
-
-    days_apart = abs((anchor_a - anchor_b).days)
-
-    if days_apart == 0:
-        return 1.0
-
-    # Logarithmic decay: 1 / (1 + log(1 + days_apart/half_life))
-    # Similar to calculate_recency_weight but for proximity between events
-    normalized_distance = days_apart / half_life_days
-    proximity = 1.0 / (1.0 + math.log1p(normalized_distance))
-
-    return proximity

hindsight_api/extensions/__init__.py

@@ -21,6 +21,10 @@ from hindsight_api.extensions.context import DefaultExtensionContext, ExtensionC
 from hindsight_api.extensions.http import HttpExtension
 from hindsight_api.extensions.loader import load_extension
 from hindsight_api.extensions.operation_validator import (
+    # Consolidation operation
+    ConsolidateContext,
+    ConsolidateResult,
+    # Core operations
     OperationValidationError,
     OperationValidatorExtension,
     RecallContext,
@@ -33,6 +37,7 @@ from hindsight_api.extensions.operation_validator import (
 )
 from hindsight_api.extensions.tenant import (
     AuthenticationError,
+    Tenant,
     TenantContext,
     TenantExtension,
 )
@@ -47,7 +52,7 @@ __all__ = [
     "DefaultExtensionContext",
     # HTTP Extension
     "HttpExtension",
-    # Operation Validator
+    # Operation Validator - Core
     "OperationValidationError",
     "OperationValidatorExtension",
     "RecallContext",
@@ -57,10 +62,14 @@ __all__ = [
     "RetainContext",
     "RetainResult",
     "ValidationResult",
+    # Operation Validator - Consolidation
+    "ConsolidateContext",
+    "ConsolidateResult",
     # Tenant/Auth
     "ApiKeyTenantExtension",
     "AuthenticationError",
     "RequestContext",
+    "Tenant",
     "TenantContext",
     "TenantExtension",
 ]

hindsight_api/extensions/builtin/tenant.py

@@ -1,6 +1,6 @@
 """Built-in tenant extension implementations."""
 
-from hindsight_api.extensions.tenant import AuthenticationError, TenantContext, TenantExtension
+from hindsight_api.extensions.tenant import AuthenticationError, Tenant, TenantContext, TenantExtension
 from hindsight_api.models import RequestContext
 
 
@@ -31,3 +31,7 @@ class ApiKeyTenantExtension(TenantExtension):
         if context.api_key != self.expected_api_key:
             raise AuthenticationError("Invalid API key")
         return TenantContext(schema_name="public")
+
+    async def list_tenants(self) -> list[Tenant]:
+        """Return public schema for single-tenant setup."""
+        return [Tenant(schema="public")]

hindsight_api/extensions/context.py

@@ -96,7 +96,7 @@ class DefaultExtensionContext(ExtensionContext):
 
     async def run_migration(self, schema: str) -> None:
         """Run migrations for a specific schema."""
-        from hindsight_api.migrations import run_migrations
+        from hindsight_api.migrations import ensure_embedding_dimension, run_migrations
 
         # Prefer getting URL from memory engine (handles pg0 case where URL is set after init)
         db_url = self._database_url
@@ -107,6 +107,15 @@ class DefaultExtensionContext(ExtensionContext):
 
         run_migrations(db_url, schema=schema)
 
+        # Ensure embedding column dimension matches the model's dimension
+        # This is needed because migrations create columns with default dimension
+        if self._memory_engine is not None:
+            embeddings = getattr(self._memory_engine, "embeddings", None)
+            if embeddings is not None:
+                dimension = getattr(embeddings, "dimension", None)
+                if dimension is not None:
+                    ensure_embedding_dimension(db_url, dimension, schema=schema)
+
     def get_memory_engine(self) -> "MemoryEngineInterface":
         """Get the memory engine interface."""
         if self._memory_engine is None:

hindsight_api/extensions/operation_validator.py

@@ -1,4 +1,4 @@
-"""Operation Validator Extension for validating retain/recall/reflect operations."""
+"""Operation Validator Extension for validating retain/recall/reflect/consolidate operations."""
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
@@ -97,6 +97,19 @@ class ReflectContext:
     context: str | None = None
 
 
+# =============================================================================
+# Consolidation Pre-operation Context
+# =============================================================================
+
+
+@dataclass
+class ConsolidateContext:
+    """Context for a consolidation operation validation (pre-operation)."""
+
+    bank_id: str
+    request_context: "RequestContext"
+
+
 # =============================================================================
 # Post-operation Contexts (includes results)
 # =============================================================================
@@ -164,9 +177,28 @@ class ReflectResultContext:
     error: str | None = None
 
 
+# =============================================================================
+# Consolidation Post-operation Context
+# =============================================================================
+
+
+@dataclass
+class ConsolidateResult:
+    """Result context for post-consolidation hook."""
+
+    bank_id: str
+    request_context: "RequestContext"
+    # Result
+    processed: int = 0
+    created: int = 0
+    updated: int = 0
+    success: bool = True
+    error: str | None = None
+
+
 class OperationValidatorExtension(Extension, ABC):
     """
-    Validates and hooks into retain/recall/reflect operations.
+    Validates and hooks into retain/recall/reflect/consolidate operations.
 
     This extension allows implementing custom logic such as:
     - Rate limiting (pre-operation)
@@ -185,9 +217,13 @@ class OperationValidatorExtension(Extension, ABC):
         -> config = {"max_requests": "100"}
 
     Hook execution order:
-        1. validate_retain/validate_recall/validate_reflect (pre-operation)
+        1. validate_* (pre-operation)
         2. [operation executes]
-        3. on_retain_complete/on_recall_complete/on_reflect_complete (post-operation)
+        3. on_*_complete (post-operation)
+
+    Supported operations:
+        - retain, recall, reflect (core memory operations)
+        - consolidate (mental models consolidation)
     """
 
     # =========================================================================
@@ -325,3 +361,44 @@ class OperationValidatorExtension(Extension, ABC):
                 - error: Error message (if failed)
         """
         pass
+
+    # =========================================================================
+    # Consolidation - Pre-operation validation hook (optional - override to implement)
+    # =========================================================================
+
+    async def validate_consolidate(self, ctx: ConsolidateContext) -> ValidationResult:
+        """
+        Validate a consolidation operation before execution.
+
+        Override to implement custom validation logic for consolidation.
+
+        Args:
+            ctx: Context containing:
+                - bank_id: Bank identifier
+                - request_context: Request context with auth info
+
+        Returns:
+            ValidationResult indicating whether the operation is allowed.
+        """
+        return ValidationResult.accept()
+
+    # =========================================================================
+    # Consolidation - Post-operation hook (optional - override to implement)
+    # =========================================================================
+
+    async def on_consolidate_complete(self, result: ConsolidateResult) -> None:
+        """
+        Called after a consolidation operation completes (success or failure).
+
+        Override to implement post-operation logic such as usage tracking or audit logging.
+
+        Args:
+            result: Result context containing:
+                - bank_id: Bank identifier
+                - processed: Number of memories processed
+                - created: Number of mental models created
+                - updated: Number of mental models updated
+                - success: Whether the operation succeeded
+                - error: Error message (if failed)
+        """
+        pass