fix

hindsight-api/hindsight_api/admin/__init__.py

@@ -0,0 +1 @@
+# Admin CLI for Hindsight

hindsight-api/hindsight_api/admin/cli.py

@@ -0,0 +1,222 @@
+"""
+Hindsight Admin CLI - backup and restore operations.
+"""
+
+import asyncio
+import io
+import json
+import logging
+import zipfile
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+import asyncpg
+import typer
+
+from ..config import HindsightConfig
+from ..pg0 import parse_pg0_url, resolve_database_url
+
+
+def _fq_table(table: str, schema: str) -> str:
+    """Get fully-qualified table name with schema prefix."""
+    return f"{schema}.{table}"
+
+
+# Setup logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(message)s",
+)
+logger = logging.getLogger(__name__)
+
+app = typer.Typer(name="hindsight-admin", help="Hindsight administrative commands")
+
+# Tables to backup/restore in dependency order
+# Import must happen in this order due to foreign key constraints
+BACKUP_TABLES = [
+    "banks",
+    "documents",
+    "entities",
+    "chunks",
+    "memory_units",
+    "unit_entities",
+    "entity_cooccurrences",
+    "memory_links",
+]
+
+MANIFEST_VERSION = "1"
+
+
+async def _backup(database_url: str, output_path: Path, schema: str = "public") -> dict[str, Any]:
+    """Backup all tables to a zip file using binary COPY protocol."""
+    conn = await asyncpg.connect(database_url)
+    try:
+        tables: dict[str, Any] = {}
+        manifest: dict[str, Any] = {
+            "version": MANIFEST_VERSION,
+            "created_at": datetime.now(timezone.utc).isoformat(),
+            "schema": schema,
+            "tables": tables,
+        }
+
+        # Use a transaction with REPEATABLE READ isolation to get a consistent
+        # snapshot across all tables. This prevents race conditions where
+        # entity_cooccurrences could reference entities created after the
+        # entities table was backed up.
+        async with conn.transaction(isolation="repeatable_read"):
+            with zipfile.ZipFile(output_path, "w", zipfile.ZIP_DEFLATED) as zf:
+                for i, table in enumerate(BACKUP_TABLES, 1):
+                    typer.echo(f"  [{i}/{len(BACKUP_TABLES)}] Backing up {table}...", nl=False)
+
+                    buffer = io.BytesIO()
+
+                    # Use binary COPY for exact type preservation
+                    # asyncpg requires schema_name as separate parameter
+                    await conn.copy_from_table(table, schema_name=schema, output=buffer, format="binary")
+
+                    data = buffer.getvalue()
+                    zf.writestr(f"{table}.bin", data)
+
+                    # Get row count for manifest
+                    qualified_table = _fq_table(table, schema)
+                    row_count = await conn.fetchval(f"SELECT COUNT(*) FROM {qualified_table}")
+                    tables[table] = {
+                        "rows": row_count,
+                        "size_bytes": len(data),
+                    }
+
+                    typer.echo(f" {row_count} rows")
+
+                zf.writestr("manifest.json", json.dumps(manifest, indent=2))
+
+        return manifest
+    finally:
+        await conn.close()
+
+
+async def _restore(database_url: str, input_path: Path, schema: str = "public") -> dict[str, Any]:
+    """Restore all tables from a zip file using binary COPY protocol."""
+    conn = await asyncpg.connect(database_url)
+    try:
+        with zipfile.ZipFile(input_path, "r") as zf:
+            # Read and validate manifest
+            manifest: dict[str, Any] = json.loads(zf.read("manifest.json"))
+            if manifest.get("version") != MANIFEST_VERSION:
+                raise ValueError(f"Unsupported backup version: {manifest.get('version')}")
+
+            # Use a transaction for atomic restore - either all tables are
+            # restored or none are, preventing partial/inconsistent state.
+            async with conn.transaction():
+                typer.echo("  Clearing existing data...")
+                # Truncate tables in reverse order (respects FK constraints)
+                for table in reversed(BACKUP_TABLES):
+                    qualified_table = _fq_table(table, schema)
+                    await conn.execute(f"TRUNCATE TABLE {qualified_table} CASCADE")
+
+                # Restore tables in forward order
+                for i, table in enumerate(BACKUP_TABLES, 1):
+                    filename = f"{table}.bin"
+                    if filename not in zf.namelist():
+                        typer.echo(f"  [{i}/{len(BACKUP_TABLES)}] {table}: skipped (not in backup)")
+                        continue
+
+                    expected_rows = manifest["tables"].get(table, {}).get("rows", "?")
+                    typer.echo(f"  [{i}/{len(BACKUP_TABLES)}] Restoring {table}... {expected_rows} rows")
+
+                    data = zf.read(filename)
+                    buffer = io.BytesIO(data)
+                    # asyncpg requires schema_name as separate parameter
+                    await conn.copy_to_table(table, schema_name=schema, source=buffer, format="binary")
+
+                # Refresh materialized view
+                typer.echo("  Refreshing materialized views...")
+                await conn.execute(f"REFRESH MATERIALIZED VIEW {_fq_table('memory_units_bm25', schema)}")
+
+        return manifest
+    finally:
+        await conn.close()
+
+
+async def _run_backup(db_url: str, output: Path, schema: str = "public") -> dict[str, Any]:
+    """Resolve database URL and run backup."""
+    is_pg0, instance_name, _ = parse_pg0_url(db_url)
+    if is_pg0:
+        typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
+    resolved_url = await resolve_database_url(db_url)
+    return await _backup(resolved_url, output, schema)
+
+
+async def _run_restore(db_url: str, input_file: Path, schema: str = "public") -> dict[str, Any]:
+    """Resolve database URL and run restore."""
+    is_pg0, instance_name, _ = parse_pg0_url(db_url)
+    if is_pg0:
+        typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
+    resolved_url = await resolve_database_url(db_url)
+    return await _restore(resolved_url, input_file, schema)
+
+
+@app.command()
+def backup(
+    output: Path = typer.Argument(..., help="Output file path (.zip)"),
+    schema: str = typer.Option("public", "--schema", "-s", help="Database schema to backup"),
+):
+    """Backup the Hindsight database to a zip file."""
+    config = HindsightConfig.from_env()
+
+    if not config.database_url:
+        typer.echo("Error: Database URL not configured.", err=True)
+        typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True)
+        raise typer.Exit(1)
+
+    if output.suffix != ".zip":
+        output = output.with_suffix(".zip")
+
+    typer.echo(f"Backing up database (schema: {schema}) to {output}...")
+
+    manifest = asyncio.run(_run_backup(config.database_url, output, schema))
+
+    total_rows = sum(t["rows"] for t in manifest["tables"].values())
+    typer.echo(f"Backed up {total_rows} rows across {len(BACKUP_TABLES)} tables")
+    typer.echo(f"Backup saved to {output}")
+
+
+@app.command()
+def restore(
+    input_file: Path = typer.Argument(..., help="Input backup file (.zip)"),
+    schema: str = typer.Option("public", "--schema", "-s", help="Database schema to restore to"),
+    yes: bool = typer.Option(False, "--yes", "-y", help="Skip confirmation prompt"),
+):
+    """Restore the database from a backup file. WARNING: This deletes all existing data."""
+    config = HindsightConfig.from_env()
+
+    if not config.database_url:
+        typer.echo("Error: Database URL not configured.", err=True)
+        typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True)
+        raise typer.Exit(1)
+
+    if not input_file.exists():
+        typer.echo(f"Error: File not found: {input_file}", err=True)
+        raise typer.Exit(1)
+
+    if not yes:
+        typer.confirm(
+            "This will DELETE all existing data and replace it with the backup. Continue?",
+            abort=True,
+        )
+
+    typer.echo(f"Restoring database (schema: {schema}) from {input_file}...")
+
+    manifest = asyncio.run(_run_restore(config.database_url, input_file, schema))
+
+    total_rows = sum(t["rows"] for t in manifest["tables"].values())
+    typer.echo(f"Restored {total_rows} rows across {len(BACKUP_TABLES)} tables")
+    typer.echo("Restore complete")
+
+
+def main():
+    app()
+
+
+if __name__ == "__main__":
+    main()

hindsight-api/hindsight_api/engine/memory_engine.py

@@ -132,7 +132,7 @@ if TYPE_CHECKING:
 
 from enum import Enum
 
-from ..pg0 import EmbeddedPostgres
+from ..pg0 import EmbeddedPostgres, parse_pg0_url
 from .entity_resolver import EntityResolver
 from .llm_wrapper import LLMConfig
 from .query_analyzer import QueryAnalyzer
@@ -259,31 +259,14 @@ class MemoryEngine(MemoryEngineInterface):
         memory_llm_base_url = memory_llm_base_url or config.get_llm_base_url() or None
         # Track pg0 instance (if used)
         self._pg0: EmbeddedPostgres | None = None
-        self._pg0_instance_name: str | None = None
 
         # Initialize PostgreSQL connection URL
         # The actual URL will be set during initialize() after starting the server
         # Supports: "pg0" (default instance), "pg0://instance-name" (named instance), or regular postgresql:// URL
-        if db_url == "pg0":
-            self._use_pg0 = True
-            self._pg0_instance_name = "hindsight"
-            self._pg0_port = None  # Use default port
-            self.db_url = None
-        elif db_url.startswith("pg0://"):
-            self._use_pg0 = True
-            # Parse instance name and optional port: pg0://instance-name or pg0://instance-name:port
-            url_part = db_url[6:]  # Remove "pg0://"
-            if ":" in url_part:
-                self._pg0_instance_name, port_str = url_part.rsplit(":", 1)
-                self._pg0_port = int(port_str)
-            else:
-                self._pg0_instance_name = url_part or "hindsight"
-                self._pg0_port = None  # Use default port
+        self._use_pg0, self._pg0_instance_name, self._pg0_port = parse_pg0_url(db_url)
+        if self._use_pg0:
             self.db_url = None
         else:
-            self._use_pg0 = False
-            self._pg0_instance_name = None
-            self._pg0_port = None
             self.db_url = db_url
 
         # Set default base URL if not provided

hindsight-api/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/pyproject.toml

@@ -38,6 +38,7 @@ dependencies = [
     "dateparser>=1.2.2",
     "google-genai>=1.0.0",
     "anthropic>=0.40.0",
+    "typer>=0.9.0",
 ]
 
 [project.optional-dependencies]
@@ -52,6 +53,7 @@ test = [
 [project.scripts]
 hindsight-api = "hindsight_api.main:main"
 hindsight-local-mcp = "hindsight_api.mcp_local:main"
+hindsight-admin = "hindsight_api.admin.cli:main"
 
 [tool.hatch.build.targets.wheel]
 packages = ["hindsight_api"]
@@ -75,7 +77,7 @@ log_cli = true
 log_cli_level = "INFO"
 log_cli_format = "%(asctime)s - %(levelname)s - %(name)s - %(message)s"
 log_cli_date_format = "%Y-%m-%d %H:%M:%S"
-addopts = "--timeout 120 -n 8 --durations=10 -v"
+addopts = "--timeout 120 -n 8 --dist loadgroup --durations=10 -v"
 asyncio_mode = "auto"
 asyncio_default_fixture_loop_scope = "function"
 log_auto_indent = true

hindsight-api/tests/test_admin_backup_restore.py

@@ -0,0 +1,208 @@
+"""
+Tests for admin backup and restore functionality.
+
+Note: These tests run sequentially (not in parallel) because they all
+manipulate the same database and do full backup/restore operations.
+"""
+
+import tempfile
+import uuid
+import zipfile
+from pathlib import Path
+
+import pytest
+
+from hindsight_api import RequestContext
+from hindsight_api.admin.cli import _backup, _restore, BACKUP_TABLES
+
+
+# Run these tests sequentially since they do full DB backup/restore
+pytestmark = pytest.mark.xdist_group(name="backup_restore")
+
+
+@pytest.mark.asyncio
+async def test_backup_restore_roundtrip(memory, pg0_db_url, request_context):
+    """Test that backup and restore preserves all data correctly."""
+    # Use unique bank ID to avoid conflicts
+    bank_id = f"test-backup-{uuid.uuid4().hex[:8]}"
+
+    # Create some test data
+    await memory.retain_batch_async(
+        bank_id=bank_id,
+        contents=[
+            {"content": "Alice is a software engineer who loves Python."},
+            {"content": "Bob works with Alice on the backend team."},
+            {"content": "The team uses PostgreSQL for their database."},
+        ],
+        request_context=request_context,
+    )
+
+    # Get counts before backup
+    async with memory._pool.acquire() as conn:
+        counts_before = {}
+        for table in BACKUP_TABLES:
+            counts_before[table] = await conn.fetchval(f"SELECT COUNT(*) FROM {table}")
+
+    # Verify we have data
+    assert counts_before["banks"] > 0
+    assert counts_before["memory_units"] > 0
+
+    # Backup to a temp file
+    with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as f:
+        backup_path = Path(f.name)
+
+    try:
+        manifest = await _backup(pg0_db_url, backup_path)
+
+        # Verify backup file exists and is valid
+        assert backup_path.exists()
+        assert backup_path.stat().st_size > 0
+
+        # Verify manifest
+        assert manifest["version"] == "1"
+        assert "created_at" in manifest
+        for table in BACKUP_TABLES:
+            assert table in manifest["tables"]
+            assert manifest["tables"][table]["rows"] == counts_before[table]
+
+        # Verify zip contents
+        with zipfile.ZipFile(backup_path, "r") as zf:
+            assert "manifest.json" in zf.namelist()
+            for table in BACKUP_TABLES:
+                assert f"{table}.bin" in zf.namelist()
+
+        # Clear all data
+        async with memory._pool.acquire() as conn:
+            for table in reversed(BACKUP_TABLES):
+                await conn.execute(f"TRUNCATE TABLE {table} CASCADE")
+
+        # Verify data is gone
+        async with memory._pool.acquire() as conn:
+            for table in BACKUP_TABLES:
+                count = await conn.fetchval(f"SELECT COUNT(*) FROM {table}")
+                assert count == 0, f"Table {table} should be empty after truncate"
+
+        # Restore from backup
+        await _restore(pg0_db_url, backup_path)
+
+        # Verify counts match original
+        async with memory._pool.acquire() as conn:
+            for table in BACKUP_TABLES:
+                count = await conn.fetchval(f"SELECT COUNT(*) FROM {table}")
+                assert count == counts_before[table], f"Table {table} count mismatch after restore"
+
+        # Verify data content is preserved
+        async with memory._pool.acquire() as conn:
+            texts = await conn.fetch(
+                "SELECT text FROM memory_units WHERE bank_id = $1",
+                bank_id,
+            )
+            text_content = " ".join(r["text"] for r in texts)
+            assert "Alice" in text_content or "software" in text_content
+
+    finally:
+        # Cleanup
+        if backup_path.exists():
+            backup_path.unlink()
+
+
+@pytest.mark.asyncio
+async def test_backup_restore_preserves_all_column_types(memory, pg0_db_url, request_context):
+    """Test that all column types are preserved: vectors, UUIDs, timestamps, JSONB."""
+    # Use unique bank ID
+    bank_id = f"test-types-{uuid.uuid4().hex[:8]}"
+
+    # Create data with meaningful content that will produce facts
+    await memory.retain_batch_async(
+        bank_id=bank_id,
+        contents=[
+            {"content": "John Smith is a senior engineer at Acme Corp since 2020."},
+            {"content": "The project deadline is December 15th 2024."},
+        ],
+        request_context=request_context,
+    )
+
+    # Get original data with all important column types
+    async with memory._pool.acquire() as conn:
+        # memory_units: UUID (id), Vector (embedding), Timestamp (event_date, created_at), JSONB (metadata)
+        original_unit = await conn.fetchrow(
+            """SELECT id, embedding, event_date, created_at, metadata, text
+               FROM memory_units WHERE bank_id = $1 LIMIT 1""",
+            bank_id,
+        )
+
+        # entities: UUID (id), Timestamp (first_seen, last_seen), JSONB (metadata)
+        original_entity = await conn.fetchrow(
+            """SELECT id, first_seen, last_seen, metadata, canonical_name
+               FROM entities WHERE bank_id = $1 LIMIT 1""",
+            bank_id,
+        )
+
+        # banks: JSONB (personality/disposition)
+        original_bank = await conn.fetchrow(
+            "SELECT bank_id, created_at, updated_at FROM banks WHERE bank_id = $1",
+            bank_id,
+        )
+
+    assert original_unit is not None, "Should have created memory units"
+    assert original_unit["embedding"] is not None, "Should have embedding"
+    assert original_unit["id"] is not None, "Should have UUID"
+    assert original_entity is not None, "Should have created entities"
+
+    with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as f:
+        backup_path = Path(f.name)
+
+    try:
+        await _backup(pg0_db_url, backup_path)
+
+        # Clear all data
+        async with memory._pool.acquire() as conn:
+            for table in reversed(BACKUP_TABLES):
+                await conn.execute(f"TRUNCATE TABLE {table} CASCADE")
+
+        await _restore(pg0_db_url, backup_path)
+
+        # Verify all column types are preserved exactly
+        async with memory._pool.acquire() as conn:
+            restored_unit = await conn.fetchrow(
+                """SELECT id, embedding, event_date, created_at, metadata, text
+                   FROM memory_units WHERE bank_id = $1 LIMIT 1""",
+                bank_id,
+            )
+
+            restored_entity = await conn.fetchrow(
+                """SELECT id, first_seen, last_seen, metadata, canonical_name
+                   FROM entities WHERE bank_id = $1 LIMIT 1""",
+                bank_id,
+            )
+
+            restored_bank = await conn.fetchrow(
+                "SELECT bank_id, created_at, updated_at FROM banks WHERE bank_id = $1",
+                bank_id,
+            )
+
+        # Verify memory_units
+        assert restored_unit is not None, "Should have restored memory unit"
+        assert restored_unit["id"] == original_unit["id"], "UUID should match exactly"
+        assert restored_unit["text"] == original_unit["text"], "Text should match"
+        assert list(restored_unit["embedding"]) == list(original_unit["embedding"]), "Vector embedding should match exactly"
+        assert restored_unit["event_date"] == original_unit["event_date"], "Timestamp should match exactly"
+        assert restored_unit["created_at"] == original_unit["created_at"], "Created timestamp should match"
+        assert restored_unit["metadata"] == original_unit["metadata"], "JSONB metadata should match"
+
+        # Verify entities
+        assert restored_entity is not None, "Should have restored entity"
+        assert restored_entity["id"] == original_entity["id"], "Entity UUID should match"
+        assert restored_entity["canonical_name"] == original_entity["canonical_name"], "Entity name should match"
+        assert restored_entity["first_seen"] == original_entity["first_seen"], "Entity first_seen should match"
+        assert restored_entity["last_seen"] == original_entity["last_seen"], "Entity last_seen should match"
+        assert restored_entity["metadata"] == original_entity["metadata"], "Entity metadata should match"
+
+        # Verify banks
+        assert restored_bank is not None, "Should have restored bank"
+        assert restored_bank["bank_id"] == original_bank["bank_id"], "Bank ID should match"
+        assert restored_bank["created_at"] == original_bank["created_at"], "Bank created_at should match"
+
+    finally:
+        if backup_path.exists():
+            backup_path.unlink()

package-lock.json

@@ -13,7 +13,7 @@
     },
     "hindsight-clients/typescript": {
       "name": "@vectorize-io/hindsight-client",
-      "version": "0.1.16",
+      "version": "0.2.1",
       "license": "MIT",
       "devDependencies": {
         "@hey-api/openapi-ts": "^0.88.0",
@@ -26,7 +26,7 @@
     },
     "hindsight-control-plane": {
       "name": "@vectorize-io/hindsight-control-plane",
-      "version": "0.1.16",
+      "version": "0.2.1",
       "license": "ISC",
       "dependencies": {
         "@radix-ui/react-alert-dialog": "^1.1.15",

uv.lock

@@ -1215,6 +1215,7 @@ dependencies = [
     { name = "tiktoken" },
     { name = "torch" },
     { name = "transformers" },
+    { name = "typer" },
     { name = "uvicorn" },
     { name = "wsproto" },
 ]
@@ -1274,6 +1275,7 @@ requires-dist = [
     { name = "tiktoken", specifier = ">=0.12.0" },
     { name = "torch", specifier = ">=2.0.0" },
     { name = "transformers", specifier = ">=4.30.0,<4.46.0" },
+    { name = "typer", specifier = ">=0.9.0" },
     { name = "uvicorn", specifier = ">=0.38.0" },
     { name = "wsproto", specifier = ">=1.0.0" },
 ]