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

hindsight_api/migrations.py

@@ -22,6 +22,7 @@ from pathlib import Path
 
 from alembic import command
 from alembic.config import Config
+from alembic.script.revision import ResolutionError
 from sqlalchemy import create_engine, text
 
 logger = logging.getLogger(__name__)
@@ -78,7 +79,18 @@ def _run_migrations_internal(database_url: str, script_location: str, schema: st
         alembic_cfg.set_main_option("target_schema", schema)
 
     # Run migrations
-    command.upgrade(alembic_cfg, "head")
+    try:
+        command.upgrade(alembic_cfg, "head")
+    except ResolutionError as e:
+        # This happens during rolling deployments when a newer version of the code
+        # has already run migrations, and this older replica doesn't have the new
+        # migration files. The database is already at a newer revision than we know.
+        # This is safe to ignore - the newer code has already applied its migrations.
+        logger.warning(
+            f"Database is at a newer migration revision than this code version knows about. "
+            f"This is expected during rolling deployments. Skipping migrations. Error: {e}"
+        )
+        return
 
     logger.info(f"Database migrations completed successfully for schema '{schema_name}'")
 
@@ -229,3 +241,131 @@ def check_migration_status(
     except Exception as e:
         logger.warning(f"Unable to check migration status: {e}")
         return None, None
+
+
+def ensure_embedding_dimension(
+    database_url: str,
+    required_dimension: int,
+    schema: str | None = None,
+) -> None:
+    """
+    Ensure the embedding column dimension matches the model's dimension.
+
+    This function checks the current vector column dimension in the database
+    and adjusts it if necessary:
+    - If dimensions match: no action needed
+    - If dimensions differ and table is empty: ALTER COLUMN to new dimension
+    - If dimensions differ and table has data: raise error with migration guidance
+
+    Args:
+        database_url: SQLAlchemy database URL
+        required_dimension: The embedding dimension required by the model
+        schema: Target PostgreSQL schema name (None for public)
+
+    Raises:
+        RuntimeError: If dimension mismatch with existing data
+    """
+    schema_name = schema or "public"
+
+    engine = create_engine(database_url)
+    with engine.connect() as conn:
+        # Check if memory_units table exists
+        table_exists = conn.execute(
+            text("""
+                SELECT EXISTS (
+                    SELECT 1 FROM information_schema.tables
+                    WHERE table_schema = :schema AND table_name = 'memory_units'
+                )
+            """),
+            {"schema": schema_name},
+        ).scalar()
+
+        if not table_exists:
+            logger.debug(f"memory_units table does not exist in schema '{schema_name}', skipping dimension check")
+            return
+
+        # Get current column dimension from pg_attribute
+        # pgvector stores dimension in atttypmod
+        current_dim = conn.execute(
+            text("""
+                SELECT atttypmod
+                FROM pg_attribute a
+                JOIN pg_class c ON a.attrelid = c.oid
+                JOIN pg_namespace n ON c.relnamespace = n.oid
+                WHERE n.nspname = :schema
+                  AND c.relname = 'memory_units'
+                  AND a.attname = 'embedding'
+            """),
+            {"schema": schema_name},
+        ).scalar()
+
+        if current_dim is None:
+            logger.warning("Could not determine current embedding dimension, skipping check")
+            return
+
+        # pgvector stores dimension directly in atttypmod (no offset like other types)
+        current_dimension = current_dim
+
+        if current_dimension == required_dimension:
+            logger.debug(f"Embedding dimension OK: {current_dimension}")
+            return
+
+        logger.info(
+            f"Embedding dimension mismatch: database has {current_dimension}, model requires {required_dimension}"
+        )
+
+        # Check if table has data
+        row_count = conn.execute(
+            text(f"SELECT COUNT(*) FROM {schema_name}.memory_units WHERE embedding IS NOT NULL")
+        ).scalar()
+
+        if row_count > 0:
+            raise RuntimeError(
+                f"Cannot change embedding dimension from {current_dimension} to {required_dimension}: "
+                f"memory_units table contains {row_count} rows with embeddings. "
+                f"To change dimensions, you must either:\n"
+                f"  1. Re-embed all data: DELETE FROM {schema_name}.memory_units; then restart\n"
+                f"  2. Use a model with {current_dimension}-dimensional embeddings"
+            )
+
+        # Table is empty, safe to alter column
+        logger.info(f"Altering embedding column dimension from {current_dimension} to {required_dimension}")
+
+        # Drop the HNSW index on embedding column if it exists
+        # Only drop indexes that use 'hnsw' and reference the 'embedding' column
+        conn.execute(
+            text(f"""
+                DO $$
+                DECLARE idx_name TEXT;
+                BEGIN
+                    FOR idx_name IN
+                        SELECT indexname FROM pg_indexes
+                        WHERE schemaname = '{schema_name}'
+                          AND tablename = 'memory_units'
+                          AND indexdef LIKE '%hnsw%'
+                          AND indexdef LIKE '%embedding%'
+                    LOOP
+                        EXECUTE 'DROP INDEX IF EXISTS {schema_name}.' || idx_name;
+                    END LOOP;
+                END $$;
+            """)
+        )
+
+        # Alter the column type
+        conn.execute(
+            text(f"ALTER TABLE {schema_name}.memory_units ALTER COLUMN embedding TYPE vector({required_dimension})")
+        )
+        conn.commit()
+
+        # Recreate the HNSW index
+        conn.execute(
+            text(f"""
+                CREATE INDEX IF NOT EXISTS idx_memory_units_embedding_hnsw
+                ON {schema_name}.memory_units
+                USING hnsw (embedding vector_cosine_ops)
+                WITH (m = 16, ef_construction = 64)
+            """)
+        )
+        conn.commit()
+
+        logger.info(f"Successfully changed embedding dimension to {required_dimension}")

hindsight_api/models.py

@@ -41,6 +41,8 @@ from sqlalchemy.dialects.postgresql import JSONB, TIMESTAMP, UUID
 from sqlalchemy.ext.asyncio import AsyncAttrs
 from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
 
+from .config import EMBEDDING_DIMENSION
+
 
 class Base(AsyncAttrs, DeclarativeBase):
     """Base class for all models."""
@@ -81,7 +83,7 @@ class MemoryUnit(Base):
     bank_id: Mapped[str] = mapped_column(Text, nullable=False)
     document_id: Mapped[str | None] = mapped_column(Text)
     text: Mapped[str] = mapped_column(Text, nullable=False)
-    embedding = mapped_column(Vector(384))  # pgvector type
+    embedding = mapped_column(Vector(EMBEDDING_DIMENSION))  # pgvector type
     context: Mapped[str | None] = mapped_column(Text)
     event_date: Mapped[datetime] = mapped_column(
         TIMESTAMP(timezone=True), nullable=False
@@ -93,7 +95,6 @@ class MemoryUnit(Base):
     mentioned_at: Mapped[datetime | None] = mapped_column(TIMESTAMP(timezone=True))  # When fact was mentioned
     fact_type: Mapped[str] = mapped_column(Text, nullable=False, server_default="world")
     confidence_score: Mapped[float | None] = mapped_column(Float)
-    access_count: Mapped[int] = mapped_column(Integer, server_default="0")
     unit_metadata: Mapped[dict] = mapped_column(
         "metadata", JSONB, server_default=sql_text("'{}'::jsonb")
     )  # User-defined metadata (str->str)
@@ -129,7 +130,6 @@ class MemoryUnit(Base):
         Index("idx_memory_units_document_id", "document_id"),
         Index("idx_memory_units_event_date", "event_date", postgresql_ops={"event_date": "DESC"}),
         Index("idx_memory_units_bank_date", "bank_id", "event_date", postgresql_ops={"event_date": "DESC"}),
-        Index("idx_memory_units_access_count", "access_count", postgresql_ops={"access_count": "DESC"}),
         Index("idx_memory_units_fact_type", "fact_type"),
         Index("idx_memory_units_bank_fact_type", "bank_id", "fact_type"),
         Index(

hindsight_api/pg0.py

@@ -132,3 +132,56 @@ async def stop_embedded_postgres() -> None:
     global _default_instance
     if _default_instance:
         await _default_instance.stop()
+
+
+def parse_pg0_url(db_url: str) -> tuple[bool, str | None, int | None]:
+    """
+    Parse a database URL and check if it's a pg0:// embedded database URL.
+
+    Supports:
+    - "pg0" -> default instance "hindsight"
+    - "pg0://instance-name" -> named instance
+    - "pg0://instance-name:port" -> named instance with explicit port
+    - Any other URL (e.g., postgresql://) -> not a pg0 URL
+
+    Args:
+        db_url: The database URL to parse
+
+    Returns:
+        Tuple of (is_pg0, instance_name, port)
+        - is_pg0: True if this is a pg0 URL
+        - instance_name: The instance name (or None if not pg0)
+        - port: The explicit port (or None for auto-assign)
+    """
+    if db_url == "pg0":
+        return True, "hindsight", None
+
+    if db_url.startswith("pg0://"):
+        url_part = db_url[6:]  # Remove "pg0://"
+        if ":" in url_part:
+            instance_name, port_str = url_part.rsplit(":", 1)
+            return True, instance_name or "hindsight", int(port_str)
+        else:
+            return True, url_part or "hindsight", None
+
+    return False, None, None
+
+
+async def resolve_database_url(db_url: str) -> str:
+    """
+    Resolve a database URL, handling pg0:// embedded database URLs.
+
+    If the URL is a pg0:// URL, starts the embedded PostgreSQL and returns
+    the actual postgresql:// connection URL. Otherwise, returns the URL unchanged.
+
+    Args:
+        db_url: Database URL (pg0://, pg0, or postgresql://)
+
+    Returns:
+        The resolved postgresql:// connection URL
+    """
+    is_pg0, instance_name, port = parse_pg0_url(db_url)
+    if is_pg0:
+        pg0 = EmbeddedPostgres(name=instance_name, port=port)
+        return await pg0.ensure_running()
+    return db_url

hindsight_api/server.py

@@ -7,6 +7,7 @@ This module provides the ASGI app for uvicorn import string usage:
 For CLI usage, use the hindsight-api command instead.
 """
 
+import logging
 import os
 import warnings
 
@@ -17,6 +18,12 @@ warnings.filterwarnings("ignore", message="websockets.server.WebSocketServerProt
 from hindsight_api import MemoryEngine
 from hindsight_api.api import create_app
 from hindsight_api.config import get_config
+from hindsight_api.extensions import (
+    DefaultExtensionContext,
+    OperationValidatorExtension,
+    TenantExtension,
+    load_extension,
+)
 
 # Disable tokenizers parallelism to avoid warnings
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -25,12 +32,42 @@ os.environ["TOKENIZERS_PARALLELISM"] = "false"
 config = get_config()
 config.configure_logging()
 
+# Load operation validator extension if configured
+operation_validator = load_extension("OPERATION_VALIDATOR", OperationValidatorExtension)
+if operation_validator:
+    logging.info(f"Loaded operation validator: {operation_validator.__class__.__name__}")
+
+# Load tenant extension if configured
+tenant_extension = load_extension("TENANT", TenantExtension)
+if tenant_extension:
+    logging.info(f"Loaded tenant extension: {tenant_extension.__class__.__name__}")
+
 # Create app at module level (required for uvicorn import string)
 # MemoryEngine reads configuration from environment variables automatically
-_memory = MemoryEngine()
+# Note: run_migrations=True by default, but migrations are idempotent so safe with workers
+_memory = MemoryEngine(
+    operation_validator=operation_validator,
+    tenant_extension=tenant_extension,
+    run_migrations=config.run_migrations_on_startup,
+)
+
+# Set extension context on tenant extension (needed for schema provisioning)
+if tenant_extension:
+    extension_context = DefaultExtensionContext(
+        database_url=config.database_url,
+        memory_engine=_memory,
+    )
+    tenant_extension.set_context(extension_context)
+    logging.info("Extension context set on tenant extension")
 
 # Create unified app with both HTTP and optionally MCP
-app = create_app(memory=_memory, http_api_enabled=True, mcp_api_enabled=config.mcp_enabled, mcp_mount_path="/mcp")
+app = create_app(
+    memory=_memory,
+    http_api_enabled=True,
+    mcp_api_enabled=config.mcp_enabled,
+    mcp_mount_path="/mcp",
+    initialize_memory=True,
+)
 
 
 if __name__ == "__main__":

hindsight_api/worker/__init__.py

@@ -0,0 +1,11 @@
+"""
+Worker package for distributed task processing.
+
+This package provides:
+- WorkerPoller: Polls PostgreSQL for pending tasks and executes them
+- main: CLI entry point for hindsight-worker
+"""
+
+from .poller import WorkerPoller
+
+__all__ = ["WorkerPoller"]

hindsight_api/worker/main.py

@@ -0,0 +1,296 @@
+"""
+Command-line interface for Hindsight Worker.
+
+Run the worker with:
+    hindsight-worker
+
+Stop with Ctrl+C (graceful shutdown).
+"""
+
+import argparse
+import asyncio
+import atexit
+import logging
+import os
+import signal
+import socket
+import sys
+import warnings
+
+from ..config import get_config
+from ..engine.task_backend import SyncTaskBackend
+from .poller import WorkerPoller
+
+# Filter deprecation warnings from third-party libraries
+warnings.filterwarnings("ignore", message="websockets.legacy is deprecated")
+warnings.filterwarnings("ignore", message="websockets.server.WebSocketServerProtocol is deprecated")
+
+# Disable tokenizers parallelism to avoid warnings
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+logger = logging.getLogger(__name__)
+
+
+def create_worker_app(poller: WorkerPoller, memory):
+    """Create a minimal FastAPI app for worker metrics and health."""
+    from fastapi import FastAPI
+    from fastapi.responses import JSONResponse, Response
+    from prometheus_client import CONTENT_TYPE_LATEST, generate_latest
+
+    from ..metrics import create_metrics_collector, get_metrics_collector, initialize_metrics
+
+    app = FastAPI(
+        title="Hindsight Worker",
+        description="Worker process for distributed task execution",
+    )
+
+    # Initialize OpenTelemetry metrics
+    try:
+        prometheus_reader = initialize_metrics(service_name="hindsight-worker", service_version="1.0.0")
+        create_metrics_collector()
+        app.state.prometheus_reader = prometheus_reader
+        logger.info("Metrics initialized - available at /metrics endpoint")
+    except Exception as e:
+        logger.warning(f"Failed to initialize metrics: {e}. Metrics will be disabled.")
+        app.state.prometheus_reader = None
+
+    # Set up DB pool metrics if available
+    metrics_collector = get_metrics_collector()
+    if memory._pool is not None and hasattr(metrics_collector, "set_db_pool"):
+        metrics_collector.set_db_pool(memory._pool)
+        logger.info("DB pool metrics configured")
+
+    @app.get(
+        "/health",
+        summary="Health check endpoint",
+        description="Returns worker health status including database connectivity",
+        tags=["Monitoring"],
+    )
+    async def health_endpoint():
+        """Health check endpoint."""
+        health = await memory.health_check()
+        health["worker_id"] = poller.worker_id
+        health["is_shutdown"] = poller.is_shutdown
+        status_code = 200 if health.get("status") == "healthy" else 503
+        return JSONResponse(content=health, status_code=status_code)
+
+    @app.get(
+        "/metrics",
+        summary="Prometheus metrics endpoint",
+        description="Exports metrics in Prometheus format for scraping",
+        tags=["Monitoring"],
+    )
+    async def metrics_endpoint():
+        """Return Prometheus metrics."""
+        metrics_data = generate_latest()
+        return Response(content=metrics_data, media_type=CONTENT_TYPE_LATEST)
+
+    @app.get(
+        "/",
+        summary="Worker info",
+        description="Basic worker information",
+        tags=["Info"],
+    )
+    async def root():
+        """Return basic worker info."""
+        return {
+            "service": "hindsight-worker",
+            "worker_id": poller.worker_id,
+            "is_shutdown": poller.is_shutdown,
+        }
+
+    return app
+
+
+def main():
+    """Main entry point for the hindsight-worker CLI."""
+    # Load configuration from environment
+    config = get_config()
+
+    parser = argparse.ArgumentParser(
+        prog="hindsight-worker",
+        description="Hindsight Worker - distributed task processor",
+    )
+
+    # Worker options
+    parser.add_argument(
+        "--worker-id",
+        default=config.worker_id or socket.gethostname(),
+        help="Worker identifier (default: hostname, env: HINDSIGHT_API_WORKER_ID)",
+    )
+    parser.add_argument(
+        "--poll-interval",
+        type=int,
+        default=config.worker_poll_interval_ms,
+        help=f"Poll interval in milliseconds (default: {config.worker_poll_interval_ms}, env: HINDSIGHT_API_WORKER_POLL_INTERVAL_MS)",
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=config.worker_batch_size,
+        help=f"Tasks to claim per poll (default: {config.worker_batch_size}, env: HINDSIGHT_API_WORKER_BATCH_SIZE)",
+    )
+    parser.add_argument(
+        "--max-retries",
+        type=int,
+        default=config.worker_max_retries,
+        help=f"Max retries before marking failed (default: {config.worker_max_retries}, env: HINDSIGHT_API_WORKER_MAX_RETRIES)",
+    )
+
+    # HTTP server options
+    parser.add_argument(
+        "--http-port",
+        type=int,
+        default=config.worker_http_port,
+        help=f"HTTP port for metrics/health endpoints (default: {config.worker_http_port}, env: HINDSIGHT_API_WORKER_HTTP_PORT)",
+    )
+    parser.add_argument(
+        "--http-host",
+        default="0.0.0.0",
+        help="HTTP host to bind (default: 0.0.0.0)",
+    )
+
+    # Logging options
+    parser.add_argument(
+        "--log-level",
+        default=config.log_level,
+        choices=["critical", "error", "warning", "info", "debug", "trace"],
+        help=f"Log level (default: {config.log_level}, env: HINDSIGHT_API_LOG_LEVEL)",
+    )
+
+    args = parser.parse_args()
+
+    # Configure logging
+    config.configure_logging()
+
+    # Import MemoryEngine here to avoid circular imports
+    from .. import MemoryEngine
+
+    print(f"Starting Hindsight Worker: {args.worker_id}")
+    print(f"  Poll interval: {args.poll_interval}ms")
+    print(f"  Batch size: {args.batch_size}")
+    print(f"  Max retries: {args.max_retries}")
+    print(f"  HTTP server: {args.http_host}:{args.http_port}")
+    print()
+
+    # Global references for cleanup
+    memory = None
+    poller = None
+
+    async def run():
+        nonlocal memory, poller
+        import uvicorn
+
+        from ..extensions import TenantExtension, load_extension
+
+        # Initialize MemoryEngine
+        # Workers use SyncTaskBackend because they execute tasks directly,
+        # they don't need to store tasks (they poll from DB)
+        memory = MemoryEngine(
+            run_migrations=False,  # Workers don't run migrations
+            task_backend=SyncTaskBackend(),
+        )
+
+        await memory.initialize()
+
+        print(f"Database connected: {config.database_url}")
+
+        # Load tenant extension for dynamic schema discovery
+        tenant_extension = load_extension("TENANT", TenantExtension)
+
+        if tenant_extension:
+            print("Tenant extension loaded - schemas will be discovered dynamically on each poll")
+        else:
+            print("No tenant extension configured, using public schema only")
+
+        # Create a single poller that handles all schemas dynamically
+        poller = WorkerPoller(
+            pool=memory._pool,
+            worker_id=args.worker_id,
+            executor=memory.execute_task,
+            poll_interval_ms=args.poll_interval,
+            batch_size=args.batch_size,
+            max_retries=args.max_retries,
+            tenant_extension=tenant_extension,
+        )
+
+        # Create the HTTP app for metrics/health
+        app = create_worker_app(poller, memory)
+
+        # Setup signal handlers for graceful shutdown
+        shutdown_requested = asyncio.Event()
+
+        def signal_handler(signum, frame):
+            print(f"\nReceived signal {signum}, initiating graceful shutdown...")
+            shutdown_requested.set()
+
+        signal.signal(signal.SIGINT, signal_handler)
+        signal.signal(signal.SIGTERM, signal_handler)
+
+        # Create uvicorn config and server
+        uvicorn_config = uvicorn.Config(
+            app,
+            host=args.http_host,
+            port=args.http_port,
+            log_level="info",  # Reduce uvicorn noise
+            access_log=False,
+        )
+        server = uvicorn.Server(uvicorn_config)
+
+        # Run the poller and HTTP server concurrently
+        poller_task = asyncio.create_task(poller.run())
+        http_task = asyncio.create_task(server.serve())
+
+        print(f"Worker started. Metrics available at http://{args.http_host}:{args.http_port}/metrics")
+
+        # Wait for shutdown signal
+        await shutdown_requested.wait()
+
+        # Graceful shutdown
+        print("Shutting down HTTP server...")
+        server.should_exit = True
+
+        print("Waiting for poller to finish...")
+        await poller.shutdown_graceful(timeout=30.0)
+        poller_task.cancel()
+        try:
+            await poller_task
+        except asyncio.CancelledError:
+            pass
+
+        # Wait for HTTP server to finish
+        try:
+            await asyncio.wait_for(http_task, timeout=5.0)
+        except asyncio.TimeoutError:
+            http_task.cancel()
+            try:
+                await http_task
+            except asyncio.CancelledError:
+                pass
+
+        # Close memory engine
+        await memory.close()
+        print("Worker shutdown complete")
+
+    def cleanup():
+        """Synchronous cleanup for atexit."""
+        if memory is not None and memory._pg0 is not None:
+            try:
+                loop = asyncio.new_event_loop()
+                loop.run_until_complete(memory._pg0.stop())
+                loop.close()
+                print("\npg0 stopped.")
+            except Exception as e:
+                print(f"\nError stopping pg0: {e}")
+
+    atexit.register(cleanup)
+
+    try:
+        asyncio.run(run())
+    except KeyboardInterrupt:
+        print("\nWorker interrupted")
+        sys.exit(0)
+
+
+if __name__ == "__main__":
+    main()