scripts/preflight.py

#!/usr/bin/env python
"""Check every external dependency this app needs before a deploy.

Each service is probed independently so a failure names the service rather than
surfacing as a generic 500 three layers up. Exits non-zero if anything required
is down, so it can gate a deploy.

    .venv/bin/python scripts/preflight.py
"""
from __future__ import annotations

import os
import sys
import time
from pathlib import Path

from dotenv import load_dotenv

PROJECT_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_ROOT))
load_dotenv(PROJECT_ROOT / ".env")

# Quiet the app's structured logging so results stay readable.
import logging

logging.disable(logging.CRITICAL)

PASS, FAIL, WARN = "PASS", "FAIL", "WARN"
results: list[tuple[str, str, str, float]] = []


def check(name: str, required: bool = True):
    """Run a probe, recording pass/fail and how long it took."""

    def decorator(fn):
        started = time.monotonic()
        try:
            detail = fn() or "ok"
            status = PASS
        except Exception as e:
            detail = f"{type(e).__name__}: {e}".replace("\n", " ")[:160]
            status = FAIL if required else WARN
        results.append((name, status, detail, time.monotonic() - started))
        return fn

    return decorator


@check("Config / env")
def _config():
    from app.core.config import settings

    missing = [
        k
        for k in ("nvidia_api_key", "database_url", "secret_key")
        if not getattr(settings, k, None)
    ]
    if missing:
        raise RuntimeError(f"missing settings: {', '.join(missing)}")
    return f"environment={settings.environment}"


@check("Database (Postgres)")
def _database():
    from sqlalchemy import create_engine, inspect, text

    engine = create_engine(os.environ["DATABASE_URL"], pool_pre_ping=True)
    with engine.connect() as c:
        version = c.execute(text("select version()")).scalar() or ""
        head = c.execute(text("select version_num from alembic_version")).scalar()
    tables = inspect(engine).get_table_names()
    if "architecture_decisions" not in tables:
        raise RuntimeError("architecture_decisions missing - run alembic upgrade head")
    return f"{version.split(',')[0][:34]} | {len(tables)} tables | head={head}"


@check("NVIDIA chat (authoring model)")
def _nvidia_chat():
    from app.core.llm_factory import get_chat_model
    from app.core.schemas import TeamRole

    llm = get_chat_model(role=TeamRole.SOLUTION_ARCHITECT, max_tokens=16)
    reply = llm.invoke([{"role": "user", "content": "Reply with the single word: ready"}])
    text = getattr(reply, "content", str(reply)).strip()
    if not text:
        raise RuntimeError("empty completion")
    return f"responded {text[:40]!r}"


@check("NVIDIA embeddings")
def _nvidia_embeddings():
    from app.core.llm_factory import get_embeddings_model

    vec = get_embeddings_model().embed_query("preflight")
    if not vec:
        raise RuntimeError("empty embedding")
    return f"dim={len(vec)}"


@check("Upstash Redis", required=False)
def _redis():
    url, token = os.getenv("UPSTASH_REDIS_REST_URL"), os.getenv("UPSTASH_REDIS_REST_TOKEN")
    if not (url and token):
        raise RuntimeError("not configured")
    import httpx

    r = httpx.get(f"{url}/set/preflight/ok", headers={"Authorization": f"Bearer {token}"}, timeout=10)
    r.raise_for_status()
    g = httpx.get(f"{url}/get/preflight", headers={"Authorization": f"Bearer {token}"}, timeout=10)
    g.raise_for_status()
    return f"set/get round-trip ok ({g.json().get('result')})"


@check("Pinecone", required=False)
def _pinecone():
    key, index = os.getenv("PINECONE_API_KEY"), os.getenv("PINECONE_INDEX")
    if not key:
        raise RuntimeError("not configured")
    from pinecone import Pinecone

    names = [i["name"] for i in Pinecone(api_key=key).list_indexes()]
    if index not in names:
        raise RuntimeError(f"index {index!r} does not exist (available: {names or 'none'})")
    return f"index {index!r} present"


@check("LangSmith tracing", required=False)
def _langsmith():
    """Verify the tracer the app actually builds can reach the ingest endpoint.

    Two false greens preceded this:

    1. `/info` is unauthenticated and returns 200 for any key, so it could not
       see a 401 at all.
    2. Probing `os.getenv("LANGSMITH_API_KEY")` against `/api/v1/sessions`
       tested *this script's* environment and the *read* path. Both were wrong.
       `load_dotenv()` at the top of this file puts the key in preflight's
       environment; the app never calls it, so the tracer's default client saw
       no key and every trace was rejected 401 by `/runs/multipart` while this
       check printed PASS. Reading also proved nothing about writing - the same
       key returned 200 from `/api/v1/sessions` and 401 from `/runs/multipart`.

    So: build the callbacks through the app's own factory, and probe the ingest
    route with the credentials that tracer ended up holding.
    """
    import httpx

    from app.core.llm_factory import get_langsmith_callbacks

    # Build the callbacks with the LangSmith/LangChain variables removed from
    # os.environ, because that is the app's real condition: `scripts/dev.sh`
    # execs uvicorn with only DATABASE_URL set, and the Docker image passes no
    # LANGSMITH_* either. `load_dotenv()` at the top of this file would
    # otherwise hand the tracer a key through a channel the app does not have,
    # which is precisely how the previous version of this check passed while
    # every trace was dropped. `settings` was populated at import time and is
    # unaffected, so a correctly-wired tracer still finds its key.
    stashed = {
        k: os.environ.pop(k)
        for k in list(os.environ)
        if k.startswith(("LANGSMITH_", "LANGCHAIN_"))
    }
    try:
        callbacks = get_langsmith_callbacks()
    finally:
        os.environ.update(stashed)

    if not callbacks:
        raise RuntimeError(
            "tracing disabled, or LANGSMITH_API_KEY missing from settings"
        )

    client = getattr(callbacks[0], "client", None)
    if client is None or not getattr(client, "api_key", None):
        raise RuntimeError(
            "tracer has no credentialed client - traces will be dropped 401"
        )

    # An empty body is rejected 422 by a *credentialed* request and 401/403 by an
    # uncredentialed one, so this distinguishes them without writing a run.
    r = httpx.post(
        f"{client.api_url}/runs/multipart",
        headers={"x-api-key": client.api_key},
        timeout=10,
    )
    if r.status_code in (401, 403):
        raise RuntimeError(
            f"{r.status_code} on /runs/multipart - traces are being dropped"
        )
    return f"ingest authenticated ({client.api_url})"


@check("Google OAuth config", required=False)
def _google():
    """Check the redirect URI is right *for this environment*.

    A localhost redirect is correct when developing locally and wrong only when
    shipping. Flagging it unconditionally meant every local run reported a
    degraded check that needed no action, which is how a warning stops being
    read. Only ENVIRONMENT=production makes localhost a real problem.
    """
    cid, secret = os.getenv("GOOGLE_CLIENT_ID"), os.getenv("GOOGLE_CLIENT_SECRET")
    if not (cid and secret):
        raise RuntimeError("not configured")

    redirect = os.getenv("GOOGLE_REDIRECT_URI", "")
    if not redirect:
        raise RuntimeError("GOOGLE_REDIRECT_URI is not set")

    is_local = "localhost" in redirect or "127.0.0.1" in redirect
    in_production = os.getenv("ENVIRONMENT", "development").lower() == "production"

    if is_local and in_production:
        raise RuntimeError(f"ENVIRONMENT=production but redirect URI is local: {redirect}")
    if is_local:
        return f"{redirect} (local - set a public URI before deploying)"
    if not redirect.startswith("https://"):
        raise RuntimeError(f"non-local redirect URI must use https: {redirect}")
    return redirect


def main() -> int:
    width = max(len(n) for n, *_ in results)
    print()
    for name, status, detail, secs in results:
        colour = {PASS: "\033[32m", FAIL: "\033[31m", WARN: "\033[33m"}[status]
        print(f"  {colour}{status:4}\033[0m  {name:<{width}}  {secs:5.2f}s  {detail}")

    failed = [n for n, s, *_ in results if s == FAIL]
    warned = [n for n, s, *_ in results if s == WARN]
    print()
    if failed:
        print(f"  \033[31m{len(failed)} required check(s) failed:\033[0m {', '.join(failed)}")
    if warned:
        print(f"  \033[33m{len(warned)} optional check(s) degraded:\033[0m {', '.join(warned)}")
    if not failed and not warned:
        print("  \033[32mAll checks passed.\033[0m")
    print()
    return 1 if failed else 0


if __name__ == "__main__":
    sys.exit(main())