start_hf.py

from __future__ import annotations

import os
import sys
from datetime import datetime, timezone
from pathlib import Path

import chromadb


ROOT = Path(__file__).resolve().parent
CHROMA_PATH = ROOT / "chroma_db_docs"
CHROMA_COLLECTION = "document_context"
MODEL_PATH = ROOT / "MODELS" / "qwen3-4b-instruct-gguf"


def log(stage: str, message: str) -> None:
    timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
    print(f"[{timestamp}] [BOOT:{stage}] {message}", flush=True)


def require_environment() -> None:
    missing = [name for name in ("DATABASE_URL",) if not os.getenv(name)]
    if missing:
        raise RuntimeError(f"Missing required Space secrets: {', '.join(missing)}")
    log("ENV", "Required secrets are present (values are not printed).")


def validate_document_chroma() -> None:
    database_path = CHROMA_PATH / "chroma.sqlite3"
    if not database_path.is_file():
        raise FileNotFoundError(f"Generated Chroma database was not found: {database_path}")

    client = chromadb.PersistentClient(path=str(CHROMA_PATH))
    collection = client.get_collection(CHROMA_COLLECTION)
    item_count = collection.count()
    if item_count == 0:
        raise RuntimeError(f"Generated Chroma collection is empty: {CHROMA_COLLECTION}")
    log("CHROMA", f"Using build-generated {CHROMA_COLLECTION} collection ({item_count} chunks).")


def validate_local_model() -> None:
    required_files = ("model.gguf",)
    missing = [name for name in required_files if not (MODEL_PATH / name).is_file()]
    if missing:
        raise FileNotFoundError(
            f"Local Qwen model is incomplete at {MODEL_PATH}; missing: {', '.join(missing)}"
        )
    log("MODEL", f"Local Qwen model is available at {MODEL_PATH}.")


def start_api() -> None:
    port = os.getenv("PORT", "7860")
    log("API", f"Starting Agent API on 0.0.0.0:{port}...")
    os.execv(
        sys.executable,
        [
            sys.executable,
            "-m",
            "uvicorn",
            "main:app",
            "--host",
            "0.0.0.0",
            "--port",
            port,
            "--log-level",
            "info",
            "--access-log",
        ],
    )


def main() -> None:
    log("START", "Starting Hugging Face Space initialization.")
    require_environment()
    validate_document_chroma()
    validate_local_model()
    start_api()


if __name__ == "__main__":
    try:
        main()
    except Exception as exc:
        log("ERROR", f"Startup aborted: {exc}")
        raise