test_embeddings.py

"""Quick test to verify ChromaDB embeddings work."""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent))

from bot.embeddings import EmbeddingManager

print("=" * 60)
print("Testing ChromaDB Embeddings")
print("=" * 60)
print()

em = EmbeddingManager()
print(f"✓ EmbeddingManager initialized")
print(f"✓ Collections exist: {em.collections_exist}")
print()

# Test the problematic column
print("Testing 'Sub Type Of The Surgery' collection...")
try:
    # Test with very low threshold to see ALL matches
    results = em.search("spine", columns=["Sub Type Of The Surgery"], top_k=10, threshold=0.0)
    print(f"✓ Search successful! Found {len(results)} results (threshold=0.0, showing top 10):")
    for i, r in enumerate(results[:10], 1):
        print(f"  {i}. {r['text'][:70]}... (similarity: {r['similarity']:.3f})")
    
    # Test with default threshold
    results_default = em.search("spine surgery", columns=["Sub Type Of The Surgery"], top_k=5)
    print(f"\n  With default threshold (0.7): {len(results_default)} results")
    if results_default:
        for r in results_default:
            print(f"    - {r['text'][:60]}... ({r['similarity']:.3f})")
except Exception as e:
    print(f"✗ Error: {e}")
    import traceback
    traceback.print_exc()

print()
print("Testing 'Descriptio Of Surgery' for 'hernia'...")
try:
    results = em.search("hernia", columns=["Descriptio Of Surgery"], top_k=10, threshold=0.0)
    print(f"✓ Found {len(results)} results (threshold=0.0, showing top 5):")
    for i, r in enumerate(results[:5], 1):
        print(f"  {i}. {r['text'][:70]}... (similarity: {r['similarity']:.3f})")
    
    results_default = em.search("hernia", columns=["Descriptio Of Surgery"], top_k=5)
    print(f"\n  With default threshold (0.7): {len(results_default)} results")
except Exception as e:
    print(f"✗ Error: {e}")

print()
print("Testing 'Reason of Admission' for 'hernia'...")
try:
    results = em.search("hernia", columns=["Reason of Admission"], top_k=5, threshold=0.0)
    print(f"✓ Found {len(results)} results (threshold=0.0, showing top 5):")
    for i, r in enumerate(results[:5], 1):
        print(f"  {i}. {r['text'][:70]}... (similarity: {r['similarity']:.3f})")
except Exception as e:
    print(f"✗ Error: {e}")

print()
print("=" * 60)
print("Test complete")
print("=" * 60)