quickstart.py

#!/usr/bin/env python3
"""
Writ Quickstart -- one-command setup and full pipeline execution.

Drop PDFs in resources/, set your credentials, run this script.

Usage:
    python quickstart.py                  # Full pipeline + tests + start server
    python quickstart.py --no-server      # Pipeline + tests, skip server
    python quickstart.py --skip-extract   # Reuse existing data/extracted_entities.json
    python quickstart.py --skip-embed     # Skip vector embedding step
    python quickstart.py --skip-tests     # Skip unit tests
"""

import argparse
import json
import os
import shutil
import subprocess
import sys
import time
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parent
RESOURCES_DIR = PROJECT_ROOT / "resources"
DATA_DIR = PROJECT_ROOT / "data"
LOGS_DIR = PROJECT_ROOT / "logs"
ENV_FILE = PROJECT_ROOT / ".env"
CONFIG_FILE = PROJECT_ROOT / "config.py"
EXTRACTED_ENTITIES = DATA_DIR / "extracted_entities.json"

REQUIRED_PYTHON = (3, 11)
NEO4J_HEALTH_URL = "http://localhost:7474"
NEO4J_BOLT_PORT = 7687
MAX_NEO4J_WAIT = 60  # seconds


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _print_step(step_num: int, total: int, message: str):
    print(f"\n{'='*60}")
    print(f"  [{step_num}/{total}] {message}")
    print(f"{'='*60}")


def _run(cmd: list[str], check: bool = True, **kwargs) -> subprocess.CompletedProcess:
    """Run a subprocess, printing the command first."""
    print(f"  $ {' '.join(cmd)}")
    return subprocess.run(cmd, check=check, **kwargs)


def _has_command(name: str) -> bool:
    return shutil.which(name) is not None


# ---------------------------------------------------------------------------
# Step 1: Preflight checks
# ---------------------------------------------------------------------------

def check_prerequisites():
    """Verify Python version and Docker availability."""
    _print_step(1, 8, "Checking prerequisites")

    v = sys.version_info
    print(f"  Python {v.major}.{v.minor}.{v.micro}")
    if (v.major, v.minor) < REQUIRED_PYTHON:
        print(f"  ERROR: Python {REQUIRED_PYTHON[0]}.{REQUIRED_PYTHON[1]}+ required")
        sys.exit(1)
    print("  Python version OK")

    if not _has_command("docker"):
        print("  WARNING: 'docker' not found — you'll need Neo4j running another way")
    else:
        print("  Docker found")

    if not _has_command("docker-compose") and not _has_command("docker"):
        print("  WARNING: 'docker-compose' not found")
    else:
        print("  Docker Compose available")


# ---------------------------------------------------------------------------
# Step 2: Create directories and config files
# ---------------------------------------------------------------------------

def setup_files():
    """Create directories, .env, and config.py from templates if missing."""
    _print_step(2, 8, "Setting up directories and config files")

    for d in (RESOURCES_DIR, DATA_DIR, LOGS_DIR):
        d.mkdir(exist_ok=True)
        print(f"  {d.relative_to(PROJECT_ROOT)}/ OK")

    if not ENV_FILE.exists():
        src = PROJECT_ROOT / "env.example"
        if src.exists():
            shutil.copy2(src, ENV_FILE)
            print("  Created .env from env.example")
            print("  >>> IMPORTANT: Edit .env and set your OPENAI_API_KEY <<<")
        else:
            print("  WARNING: env.example not found, cannot create .env")
    else:
        print("  .env already exists")

    if not CONFIG_FILE.exists():
        src = PROJECT_ROOT / "config.example.py"
        if src.exists():
            shutil.copy2(src, CONFIG_FILE)
            print("  Created config.py from config.example.py")
        else:
            print("  WARNING: config.example.py not found")
    else:
        print("  config.py already exists")


# ---------------------------------------------------------------------------
# Step 3: Install dependencies
# ---------------------------------------------------------------------------

def install_dependencies():
    """Install Python dependencies from pyproject.toml."""
    _print_step(3, 8, "Installing Python dependencies")
    _run([sys.executable, "-m", "pip", "install", "-e", ".[dev]", "--quiet"])
    print("  Dependencies installed")


# ---------------------------------------------------------------------------
# Step 4: Start Neo4j
# ---------------------------------------------------------------------------

def start_neo4j():
    """Start Neo4j via docker-compose and wait until healthy."""
    _print_step(4, 8, "Starting Neo4j (Docker)")

    compose_file = PROJECT_ROOT / "docker-compose.yml"
    if not compose_file.exists():
        print("  No docker-compose.yml found, assuming Neo4j is running externally")
        return

    # Determine compose command (docker compose vs docker-compose)
    if _has_command("docker-compose"):
        compose_cmd = ["docker-compose"]
    elif _has_command("docker"):
        compose_cmd = ["docker", "compose"]
    else:
        print("  Docker not available — ensure Neo4j is running at bolt://localhost:7687")
        return

    _run([*compose_cmd, "up", "-d"], check=False)
    print(f"  Waiting up to {MAX_NEO4J_WAIT}s for Neo4j to be ready...")

    import socket
    start = time.time()
    while time.time() - start < MAX_NEO4J_WAIT:
        try:
            with socket.create_connection(("localhost", NEO4J_BOLT_PORT), timeout=2):
                # Connection accepted — now verify Neo4j is actually ready
                time.sleep(2)
                print(f"  Neo4j is ready ({int(time.time() - start)}s)")
                return
        except (ConnectionRefusedError, OSError, socket.timeout):
            time.sleep(2)

    print(f"  WARNING: Neo4j may not be ready after {MAX_NEO4J_WAIT}s — continuing anyway")


# ---------------------------------------------------------------------------
# Step 5: Validate environment
# ---------------------------------------------------------------------------

def validate_env() -> dict:
    """Load .env and check required credentials are set."""
    _print_step(5, 8, "Validating environment")

    # Load .env into os.environ
    try:
        from dotenv import load_dotenv
        load_dotenv(ENV_FILE, override=False)
    except ImportError:
        # Fallback: parse .env manually
        if ENV_FILE.exists():
            for line in ENV_FILE.read_text().splitlines():
                line = line.strip()
                if line and not line.startswith("#") and "=" in line:
                    key, _, value = line.partition("=")
                    os.environ.setdefault(key.strip(), value.strip())

    env = {
        "NEO4J_URI": os.environ.get("NEO4J_URI", "bolt://localhost:7687"),
        "NEO4J_USER": os.environ.get("NEO4J_USER", "neo4j"),
        "NEO4J_PASSWORD": os.environ.get("NEO4J_PASSWORD", ""),
        "OPENAI_API_KEY": os.environ.get("OPENAI_API_KEY", ""),
    }

    if not env["NEO4J_PASSWORD"]:
        print("  ERROR: NEO4J_PASSWORD not set in .env")
        sys.exit(1)

    if not env["OPENAI_API_KEY"]:
        print("  ERROR: OPENAI_API_KEY not set in .env")
        print("  Edit .env and add your OpenAI API key, then re-run.")
        sys.exit(1)

    # Check for PDFs
    pdfs = list(RESOURCES_DIR.glob("*.pdf"))
    print(f"  Neo4j URI:    {env['NEO4J_URI']}")
    print(f"  OpenAI key:   {'***' + env['OPENAI_API_KEY'][-4:] if len(env['OPENAI_API_KEY']) > 4 else '(set)'}")
    print(f"  PDFs found:   {len(pdfs)} in resources/")

    if not pdfs:
        print("  WARNING: No PDFs in resources/ — extraction will produce no data.")
        print("  Drop your regulatory PDFs in resources/ and re-run, or press Enter to continue.")
        try:
            input("  Press Enter to continue anyway, or Ctrl+C to abort... ")
        except KeyboardInterrupt:
            print("\n  Aborted.")
            sys.exit(0)

    return env


# ---------------------------------------------------------------------------
# Step 6: Run pipeline
# ---------------------------------------------------------------------------

def run_extraction(skip: bool = False):
    """Extract entities from PDFs using LLM."""
    _print_step(6, 8, "Entity extraction")

    if skip and EXTRACTED_ENTITIES.exists():
        with open(EXTRACTED_ENTITIES) as f:
            data = json.load(f)
        print(f"  Skipped — using existing {EXTRACTED_ENTITIES.name}")
        print(f"  ({data.get('entity_count', '?')} entities, {data.get('relationship_count', '?')} relationships)")
        return

    pdfs = list(RESOURCES_DIR.glob("*.pdf"))
    if not pdfs:
        print("  No PDFs in resources/ — skipping extraction")
        return

    _run([
        sys.executable,
        str(PROJECT_ROOT / "scripts" / "entity_extractor_template.py"),
    ])
    print("  Extraction complete")


def run_ingestion():
    """Ingest extracted entities into Neo4j."""
    if not EXTRACTED_ENTITIES.exists():
        print("  No extracted_entities.json — skipping ingestion")
        return

    print("  Ingesting into Neo4j...")
    _run([
        sys.executable,
        str(PROJECT_ROOT / "scripts" / "ingest_template.py"),
    ])
    print("  Ingestion complete")


def run_embedding(skip: bool = False):
    """Create vector embeddings for RAG."""
    if skip:
        print("  Skipped embedding step")
        return

    pdfs = list(RESOURCES_DIR.glob("*.pdf"))
    if not pdfs:
        print("  No PDFs — skipping embedding")
        return

    print("  Creating vector embeddings...")
    _run([
        sys.executable,
        str(PROJECT_ROOT / "scripts" / "embed_documents.py"),
    ])
    print("  Embedding complete")


def run_pipeline(skip_extract: bool = False, skip_embed: bool = False):
    """Run the full data pipeline: extract -> ingest -> embed."""
    run_extraction(skip=skip_extract)

    _print_step(7, 8, "Ingestion + Embedding")
    run_ingestion()
    run_embedding(skip=skip_embed)


# ---------------------------------------------------------------------------
# Step 7 (interleaved): Run tests
# ---------------------------------------------------------------------------

def run_tests(skip: bool = False):
    """Run unit tests to verify installation."""
    if skip:
        return

    print("\n  Running unit tests...")
    result = _run(
        [sys.executable, "-m", "pytest", "tests/", "-m", "unit", "-v", "--tb=short"],
        check=False,
        cwd=str(PROJECT_ROOT),
    )
    if result.returncode == 0:
        print("  All unit tests passed")
    else:
        print("  WARNING: Some tests failed (see output above)")
        print("  Pipeline will continue — check test output for details")


# ---------------------------------------------------------------------------
# Step 8: Start server
# ---------------------------------------------------------------------------

def start_server():
    """Start the local development server."""
    _print_step(8, 8, "Starting local server")
    print("  Open http://localhost:8080 in your browser")
    print("  Press Ctrl+C to stop\n")

    _run([sys.executable, str(PROJECT_ROOT / "server.py")])


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="Writ quickstart — one-command setup and pipeline",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            "Steps performed:\n"
            "  1. Check prerequisites (Python, Docker)\n"
            "  2. Create directories and config files\n"
            "  3. Install Python dependencies\n"
            "  4. Start Neo4j via Docker\n"
            "  5. Validate environment (.env credentials, PDFs)\n"
            "  6. Extract entities from PDFs (LLM)\n"
            "  7. Ingest into Neo4j + create embeddings\n"
            "  8. Run unit tests + start server\n"
        ),
    )
    parser.add_argument("--no-server", action="store_true", help="Do everything except start the server")
    parser.add_argument("--skip-extract", action="store_true", help="Reuse existing extracted_entities.json")
    parser.add_argument("--skip-embed", action="store_true", help="Skip vector embedding step")
    parser.add_argument("--skip-tests", action="store_true", help="Skip unit tests")
    parser.add_argument("--skip-install", action="store_true", help="Skip pip install (deps already installed)")
    parser.add_argument("--skip-docker", action="store_true", help="Skip Docker/Neo4j startup (already running)")
    args = parser.parse_args()

    print("=" * 60)
    print("  WRIT QUICKSTART")
    print("  Transparent AI knowledge graphs from regulatory documents")
    print("=" * 60)

    # 1. Prerequisites
    check_prerequisites()

    # 2. Files and directories
    setup_files()

    # 3. Install dependencies
    if not args.skip_install:
        install_dependencies()
    else:
        print("\n  Skipping dependency installation (--skip-install)")

    # 4. Start Neo4j
    if not args.skip_docker:
        start_neo4j()
    else:
        print("\n  Skipping Docker startup (--skip-docker)")

    # 5. Validate environment
    validate_env()

    # 6-7. Pipeline
    run_pipeline(
        skip_extract=args.skip_extract,
        skip_embed=args.skip_embed,
    )

    # Tests
    run_tests(skip=args.skip_tests)

    # 8. Server
    if args.no_server:
        print("\n" + "=" * 60)
        print("  QUICKSTART COMPLETE (--no-server)")
        print("  To start the server later: python server.py")
        print("=" * 60)
    else:
        start_server()


if __name__ == "__main__":
    main()