scripts/mintoak/test_db_migration.py

import os
import sys
import json

# Add project root to path
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from chromadb.utils import embedding_functions
from vector_db import get_vector_db

def run_tests():
    print(f"=== Testing Vector DB Migration ===")
    db_type = os.getenv("VECTOR_DB_TYPE", "chroma").lower()
    print(f"Active DB Type: {db_type}")
    
    _emb_fn = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="all-MiniLM-L6-v2")
    
    try:
        db = get_vector_db(_emb_fn)
    except Exception as e:
        print(f"❌ Failed to initialize DB client: {e}")
        return False
        
    # Sample test data
    test_ids = ["test_id_1", "test_id_2"]
    test_docs = [
        "Title: SmartPayments\nContent: Mintoak SmartPayments is a modular, cloud-native payments suite.",
        "Title: DigiOnboard\nContent: Mintoak DigiOnboard digitizes merchant acquisition and KYC."
    ]
    test_metadatas = [
        {"url": "https://www.mintoak.com/smartpayments", "title": "SmartPayments", "category": "Product offering"},
        {"url": "https://www.mintoak.com/digionboard", "title": "DigiOnboard", "category": "Product offering"}
    ]
    
    print("\n1. Generating test embeddings...")
    test_embs = _emb_fn(test_docs)
    print(f"Generated {len(test_embs)} embeddings of size {len(test_embs[0])}")
    
    print("\n2. Inserting test documents...")
    try:
        db.add_documents(test_ids, test_docs, test_metadatas, test_embs)
        print("✅ Documents inserted successfully.")
    except Exception as e:
        print(f"❌ Insertion failed: {e}")
        return False
        
    print("\n3. Testing count...")
    try:
        count = db.count()
        print(f"✅ DB Count: {count}")
    except Exception as e:
        print(f"❌ Count failed: {e}")
        return False
        
    print("\n4. Testing query (Retrieval)...")
    try:
        query_text = "smartpayments"
        query_emb = _emb_fn([query_text])[0]
        results = db.query(query_emb, n_results=1)
        print(f"✅ Query results: {results}")
        
        # Verify retrieved data
        if results and results.get("documents") and results["documents"][0]:
            doc = results["documents"][0][0]
            print(f"✅ Retrieved Document: '{doc}'")
            if "SmartPayments" in doc:
                print("✅ Retrieval match verification: SUCCESS")
            else:
                print("❌ Retrieval match verification: FAILED")
        else:
            print("❌ Query returned no documents.")
            return False
    except Exception as e:
        print(f"❌ Query failed: {e}")
        return False
        
    print("\n5. Testing get by ID...")
    try:
        results = db.get(["test_id_2"])
        print(f"✅ Get results: {results}")
        if results and results.get("documents") and len(results["documents"]) > 0:
            doc = results["documents"][0]
            print(f"✅ Retrieved Doc by ID: '{doc}'")
            if "DigiOnboard" in doc:
                print("✅ Get by ID verification: SUCCESS")
            else:
                print("❌ Get by ID verification: FAILED")
        else:
            print("❌ Get by ID returned no documents.")
            return False
    except Exception as e:
        print(f"❌ Get by ID failed: {e}")
        return False

    print("\n6. Testing query with metadata filtering...")
    try:
        query_text = "onboard"
        query_emb = _emb_fn([query_text])[0]
        # Filter that should NOT match our test documents (which have category "Product offering")
        filter_dict = {"category": "Company Info"}
        results = db.query(query_emb, n_results=1, where=filter_dict)
        print(f"✅ Query with filter results: {results}")
        if results and results.get("metadatas") and len(results["metadatas"][0]) > 0:
            category = results["metadatas"][0][0].get("category")
            if category == "Company Info":
                print("✅ Query filter verification: SUCCESS (correctly matched category filter)")
            else:
                print(f"❌ Query filter: returned category '{category}' which does not match filter 'Company Info'")
                return False
        else:
            print("✅ Query filter verification: SUCCESS (no matching documents found, which is valid if DB is clean)")
    except Exception as e:
        print(f"❌ Query with filter failed: {e}")
        return False

    print("\n🎉 ALL TESTS COMPLETED SUCCESSFULLY!")
    return True

if __name__ == "__main__":
    success = run_tests()
    sys.exit(0 if success else 1)