scripts/verify_huggingface_inference.py

#!/usr/bin/env python3
"""Smoke-test Hugging Face Inference API (embeddings + chat) using your .env.

Run from the repository root (so `.env` is found):

  LLM_PROVIDER=huggingface uv run python scripts/verify_huggingface_inference.py

Or set `LLM_PROVIDER=huggingface` in `.env` together with `HUGGINGFACE_API_KEY`, then:

  uv run python scripts/verify_huggingface_inference.py
"""

from __future__ import annotations

import os
import sys
from pathlib import Path

# Ensure repo root is cwd so pydantic-settings loads `.env`.
_ROOT = Path(__file__).resolve().parent.parent
os.chdir(_ROOT)
if str(_ROOT) not in sys.path:
    sys.path.insert(0, str(_ROOT))

from langchain_core.messages import HumanMessage

from api.config import get_settings
from rag.embedder import create_embedding_function
from rag.retriever import _create_chat_model


def main() -> int:
    get_settings.cache_clear()
    settings = get_settings()

    if settings.llm_provider.lower() != "huggingface":
        print(
            "Set LLM_PROVIDER=huggingface in your environment or `.env` before running this script.\n"
            f"Current LLM_PROVIDER={settings.llm_provider!r}",
            file=sys.stderr,
        )
        return 1

    if not (settings.huggingface_api_key or "").strip():
        print(
            "Missing token: set HUGGINGFACE_API_KEY or HF_TOKEN in `.env`.\n"
            "Current LLM_PROVIDER is huggingface but no Hugging Face token resolved.",
            file=sys.stderr,
        )
        return 1

    print("Model:", settings.huggingface_model)
    print("Embedding model:", settings.huggingface_embedding_model)
    print("--- Embeddings ---")
    emb = create_embedding_function()
    vec = emb.embed_query("hello from DocuAudit local smoke test")
    print(f"OK: single query embedding length = {len(vec)}")

    print("--- Chat (Inference API) ---")
    llm = _create_chat_model(settings)
    msg = HumanMessage(content='Reply with exactly the word "ok" and nothing else.')
    out = llm.invoke([msg])
    text = (getattr(out, "content", None) or str(out)).strip()
    print("OK: chat response (first 200 chars):", text[:200])

    print("\nHugging Face path looks good for deployment smoke checks.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())