test_search.py

#!/usr/bin/env python3
"""
Test script to verify semantic search is working after fixing embeddings
"""

import sys
import os
sys.path.append('.')

# Mock the missing gradio module to avoid import errors
class MockGradio:
    def __getattr__(self, name):
        return lambda *args, **kwargs: None
sys.modules['gradio'] = MockGradio()

from note_graph import NoteGraph

def main():
    print("🔍 Testing semantic search functionality...")

    try:
        # Initialize NoteGraph
        ng = NoteGraph()
        print('✅ NoteGraph initialized successfully')

        # Initialize embedding model for sentence transformers
        print('🧠 Initializing sentence transformer model...')
        ng.init_embedding_model(use_openai=False)
        print('✅ Embedding model initialized')

        # Test semantic search
        print('\n🔎 Testing semantic search for "AI"...')
        results = ng.semantic_search('AI', limit=3)
        print(f'🎯 Found {len(results)} results')

        if results:
            print('\n📋 Search Results:')
            for i, result in enumerate(results, 1):
                title = result.get('title', 'Unknown')
                score = result.get('score', 0)
                print(f'{i}. {title} (similarity: {score:.3f})')
        else:
            print('❌ No results found')

        # Test another search
        print('\n🔎 Testing semantic search for "cooking"...')
        results = ng.semantic_search('cooking', limit=2)
        print(f'🎯 Found {len(results)} results')

        if results:
            print('\n📋 Search Results:')
            for i, result in enumerate(results, 1):
                title = result.get('title', 'Unknown')
                score = result.get('score', 0)
                print(f'{i}. {title} (similarity: {score:.3f})')

        print('\n✅ Semantic search is working correctly!')
        return True

    except Exception as e:
        print(f'❌ Error testing search: {e}')
        import traceback
        traceback.print_exc()
        return False

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