preflight.py

#!/usr/bin/env python3
"""
preflight.py — Check that all services are reachable and working.

Usage:
    python preflight.py          # Run all checks
    python preflight.py llm      # Only LLM
    python preflight.py embedding # Only embedding
    python preflight.py books    # Only Google Books + Open Library

Exit codes:
    0 = all passed
    1 = at least one failed
"""

import os
import sys
from pathlib import Path

try:
  from dotenv import load_dotenv
  load_dotenv(Path(__file__).parent / ".env")
except ImportError:
  pass

GREEN = "\033[32m"
RED = "\033[31m"
CYAN = "\033[36m"
RESET = "\033[0m"

passed = 0
failed = 0


def check(name, ok, msg):
  global passed, failed
  if ok:
    print(f"  [{GREEN}OK{RESET}] {name}: {msg}")
    passed += 1
  else:
    print(f"  [{RED}FAIL{RESET}] {name}: {msg}")
    failed += 1


def skip(name, msg):
  print(f"  [{CYAN}SKIP{RESET}] {name}: {msg}")


# ============================================================
# LLM
# ============================================================

def test_llm():
  print(f"\n{CYAN}--- LLM connectivity + response ---{RESET}")

  provider = os.environ.get("LLM_PROVIDER", "").strip()
  if not provider:
    skip("LLM", "LLM_PROVIDER not set")
    return

  if provider != "openai_llm":
    skip("LLM", f"unknown provider: {provider}")
    return

  base_url = os.environ.get("LLM_BASE_URL", "")
  api_key = os.environ.get("LLM_API_KEY", "")
  model = os.environ.get("LLM_MODEL", "")
  max_tokens = int(os.environ.get("LLM_MAX_TOKENS", "256"))
  timeout = int(os.environ.get("LLM_TIMEOUT", "120"))

  if not base_url or not model:
    check("LLM config", False, "LLM_BASE_URL or LLM_MODEL not set")
    return

  import openai_llm

  # Test 1: obvious match
  idx, reason = openai_llm.pick_best_match(
      "The Hobbit", "J.R.R. Tolkien",
      [
          {"title": "The Hobbit", "author": "J. R. R. Tolkien",
              "narrator": "Andy Serkis"},
          {"title": "The Hobbit (Dramatised)",
           "author": "J.R.R. Tolkien", "narrator": "Heron Carvic"},
          {"title": "The Hobbit", "author": "J. R. R. Tolkien",
              "narrator": "Rob Inglis"},
          {"title": "The Hobbit", "author": "J. R. R. Tolkien",
              "narrator": "Martin Shaw"},
      ],
      base_url=base_url, api_key=api_key, model=model, max_tokens=max_tokens, timeout=timeout,
  )

  if reason.startswith("connection_error"):
    check("LLM reachable", False, reason)
    return
  if reason.startswith("http_"):
    check("LLM API", False, reason)
    return
  if reason.startswith("empty_response"):
    check("LLM response", False, reason)
    return
  if reason.startswith("invalid_response"):
    check("LLM parse", False, reason)
    return

  # Any non-dramatised pick is fine (0, 2, 3)
  test1_ok = idx in (0, 2, 3) and reason == "matched"
  check("LLM test 1 (ambiguous editions)", test1_ok,
        f"picked [{idx}] reason={reason}" + (" (expected 0,2,3)" if not test1_ok else ""))

  # Test 2: no match
  idx2, reason2 = openai_llm.pick_best_match(
      "The Silmarillion", "J.R.R. Tolkien",
      [
          {"title": "Dune", "author": "Frank Herbert"},
          {"title": "Harry Potter", "author": "J.K. Rowling"},
          {"title": "Project Hail Mary", "author": "Andy Weir"},
      ],
      base_url=base_url, api_key=api_key, model=model, max_tokens=max_tokens, timeout=timeout,
  )
  test2_ok = idx2 is None and reason2 == "none"
  check("LLM test 2 (no match)", test2_ok,
        f"idx={idx2} reason={reason2}" + (" (expected NONE)" if not test2_ok else ""))

  # Test 3: long messy title with series disambiguation
  idx3, reason3 = openai_llm.pick_best_match(
      "Percy Jackson 01 - The Lightning Thief - Percy Jackson and the Olympians Series, Book 1",
      "Rick Riordan",
      [
          {"title": "Percy Jackson and the Lightning Thief", "author": "Rick Riordan",
           "narrator": "Jesse Bernstein", "series": "Percy Jackson and the Olympians", "series_sequence": "1"},
          {"title": "The Sea of Monsters", "author": "Rick Riordan",
           "narrator": "Jesse Bernstein", "series": "Percy Jackson and the Olympians", "series_sequence": "2"},
          {"title": "The Titans Curse", "author": "Rick Riordan",
           "narrator": "Jesse Bernstein", "series": "Percy Jackson and the Olympians", "series_sequence": "3"},
          {"title": "The Red Pyramid", "author": "Rick Riordan",
           "narrator": "Kevin R. Free", "series": "The Kane Chronicles", "series_sequence": "1"},
          {"title": "The Lost Hero", "author": "Rick Riordan",
           "narrator": "Joshua Swanson", "series": "The Heroes of Olympus", "series_sequence": "1"},
      ],
      base_url=base_url, api_key=api_key, model=model, max_tokens=max_tokens, timeout=timeout,
  )
  test3_ok = idx3 == 0 and reason3 == "matched"
  check("LLM test 3 (long title + series)", test3_ok,
        f"idx={idx3} reason={reason3}" + (" (expected 0)" if not test3_ok else ""))


# ============================================================
# Embedding
# ============================================================

def test_embedding():
  print(f"\n{CYAN}--- Embedding model ---{RESET}")

  model_name = os.environ.get("EMBEDDING_MODEL", "").strip()
  if not model_name:
    skip("Embedding", "EMBEDDING_MODEL not set")
    return

  try:
    import embedding
  except ImportError:
    check("Embedding import", False, "sentence-transformers not installed")
    return

  device = os.environ.get("EMBEDDING_DEVICE", "")

  idx, score = embedding.pick_best_match(
      "The Hobbit", "J.R.R. Tolkien",
      [
          {"title": "The Hobbit", "author": "J.R.R. Tolkien"},
          {"title": "Harry Potter", "author": "J.K. Rowling"},
      ],
      model_name=model_name, threshold=0.5, device=device,
  )

  ok = idx == 0 and score > 0.8
  check("Embedding match", ok,
        f"picked [{idx}] score={score:.3f} device={embedding._device}"
        + ("" if ok else " (expected idx=0, score>0.8)"))


# ============================================================
# Book APIs
# ============================================================

def test_books():
  print(f"\n{CYAN}--- Google Books API ---{RESET}")
  try:
    from bookinfo import search_google_books
    results = search_google_books("The Hobbit", "Tolkien", max_results=1)
    if results and "hobbit" in results[0].get("title", "").lower():
      check("Google Books", True, results[0]["title"])
    elif results:
      check("Google Books", False,
            f"got: {results[0].get('title', '?')}")
    else:
      check("Google Books", False,
            "no results (rate limited or unreachable)")
  except Exception as e:
    check("Google Books", False, str(e))

  print(f"\n{CYAN}--- Open Library API ---{RESET}")
  try:
    from bookinfo import search_open_library
    results = search_open_library("The Hobbit", "Tolkien", max_results=1)
    if results and "hobbit" in results[0].get("title", "").lower():
      check("Open Library", True, results[0]["title"])
    elif results:
      check("Open Library", False,
            f"got: {results[0].get('title', '?')}")
    else:
      check("Open Library", False, "no results or unreachable")
  except Exception as e:
    check("Open Library", False, str(e))


# ============================================================
# Main
# ============================================================

if __name__ == "__main__":
  suite = sys.argv[1] if len(sys.argv) > 1 else "all"

  if suite in ("all", "llm"):
    test_llm()
  if suite in ("all", "embedding"):
    test_embedding()
  if suite in ("all", "books"):
    test_books()

  print(f"\n{'='*50}")
  if failed == 0:
    print(f"  {GREEN}All {passed} checks passed{RESET}")
  else:
    print(f"  {RED}{failed} failed{RESET} / {passed + failed} total")
  print(f"{'='*50}")

  sys.exit(1 if failed else 0)