scripts/test_embedding_singleton.py

"""
Smoke test: verify embedding model singleton is resilient to module-reload scenarios.

Simulates the exact failure mode reported in the evaluation:
1. First call to _get_embedding_model() succeeds (registers + instantiates)
2. Resetting the module-level global to None (simulating a module reload / failed-first-attempt)
3. Second call should NOT crash with "already registered" — the fix must handle
   fastembed's class-level registry being pre-populated.

Usage:
    cd AI Chatbot
    python -m scripts.test_embedding_singleton
"""

import sys
import time

sys.stdout.reconfigure(encoding="utf-8")  # Windows-safe emoji

print("=" * 60)
print("Smoke Test: Embedding Singleton Resilience")
print("=" * 60)


def step(label: str, ok: bool):
    status = "PASS" if ok else "FAIL"
    print(f"  [{status}] {label}")


# ---- Step 1: first-time load (should succeed) ----
print("\n[Step 1] First-time load...")
start = time.perf_counter()
from app.vector_store import _get_embedding_model, embed_text

try:
    model = _get_embedding_model()
    step("_get_embedding_model() returned an instance", model is not None)
except Exception as e:
    step(f"_get_embedding_model() raised: {e}", False)
    sys.exit(1)

elapsed = time.perf_counter() - start
print(f"  (took {elapsed:.1f}s — includes model download if first run)\n")


# ---- Step 2: second call with cached model (should return immediately) ----
print("[Step 2] Second call (cached)...")
try:
    model2 = _get_embedding_model()
    step("_get_embedding_model() returned the same instance", model2 is model)
except Exception as e:
    step(f"_get_embedding_model() raised: {e}", False)
    sys.exit(1)
print()


# ---- Step 3: embed a test query (should produce 1024-dim vector) ----
print("[Step 3] embed_text produces a valid 1024-dim vector...")
try:
    vec = embed_text("berapa pH optimal untuk selada hidroponik?")
    is_1024 = len(vec) == 1024
    has_values = any(abs(v) > 1e-8 for v in vec[:10])
    step(f"embed_text returned {len(vec)} floats", is_1024)
    step("first 10 values are non-zero (valid embedding)", has_values)
except Exception as e:
    step(f"embed_text raised: {e}", False)
    sys.exit(1)
print()


# ---- Step 4: simulate module reload (reset global, keep class registry) ----
print("[Step 4] Simulating module reload — reset _embedding_model to None...")
import app.vector_store as vs

try:
    old_model = vs._embedding_model
    vs._embedding_model = None
    step("_embedding_model reset to None", vs._embedding_model is None)
except Exception as e:
    step(f"reset failed: {e}", False)
    sys.exit(1)
print()


# ---- Step 5: the critical test — _get_embedding_model() with registry populated ----
print("[Step 5] Calling _get_embedding_model() after 'module reload'...")
print("  (fastembed class-level registry already has the model from Step 1)")
try:
    model3 = _get_embedding_model()
    step("_get_embedding_model() did NOT crash with 'already registered'", model3 is not None)
    step("instance is valid", hasattr(model3, "query_embed"))
except Exception as e:
    step(f"CRASHED: {e}", False)
    sys.exit(1)
print()


# ---- Step 6: verify embed still works after reload simulation ----
print("[Step 6] embed_text works after 'module reload'...")
try:
    vec2 = embed_text("bagaimana cara mengatasi defisiensi kalium pada tanaman cabai?")
    is_1024 = len(vec2) == 1024
    step(f"embed_text returned {len(vec2)} floats", is_1024)
    # Quick sanity: same query should produce similar embedding
    from app.vector_store import EMBEDDING_DIM
    step(f"matches EMBEDDING_DIM={EMBEDDING_DIM}", len(vec2) == EMBEDDING_DIM)
except Exception as e:
    step(f"embed_text raised: {e}", False)
    sys.exit(1)
print()


# ---- Done ----
print("=" * 60)
print("ALL STEPS PASSED — singleton fix is working correctly")
print("=" * 60)