test_note_graph.py

#!/usr/bin/env python3
"""
Comprehensive test script for NoteGraph functionality.
Tests all major components: database, embeddings, search, and knowledge graph.
"""

import os
import sys
import json
import tempfile
from note_graph import NoteGraph

def test_database_operations():
    """Test basic database operations."""
    print("🧪 Testing database operations...")

    # Use temporary database for testing
    with tempfile.NamedTemporaryFile(suffix='.db', delete=False) as tmp:
        db_path = tmp.name

    try:
        ng = NoteGraph(db_path)
        ng.init_embedding_model(use_openai=True)

        # Test adding notes
        note1_id = ng.add_note(
            "Machine Learning Basics",
            "Machine learning is a subset of artificial intelligence that enables computers to learn from data without being explicitly programmed.",
            ["AI", "ML", "Basics"]
        )
        print(f"✅ Added note 1 with ID: {note1_id}")

        note2_id = ng.add_note(
            "Neural Networks",
            "Neural networks are computing systems inspired by biological neural networks that constitute animal brains.",
            ["AI", "Neural Networks", "Deep Learning"]
        )
        print(f"✅ Added note 2 with ID: {note2_id}")

        note3_id = ng.add_note(
            "Python Programming",
            "Python is a high-level programming language known for its simplicity and readability.",
            ["Programming", "Python", "Software Development"]
        )
        print(f"✅ Added note 3 with ID: {note3_id}")

        # Test retrieving notes
        note = ng.get_note_with_connections(note1_id)
        assert note is not None
        assert note['title'] == "Machine Learning Basics"
        assert "AI" in note['tags']
        print("✅ Note retrieval successful")

        # Test text search
        conn = ng._get_connection()
        try:
            cursor = conn.execute(
                "SELECT COUNT(*) FROM notes WHERE title LIKE ? OR content LIKE ?",
                ("%machine%", "%learning%")
            )
            count = cursor.fetchone()[0]
        finally:
            conn.close()
        assert count > 0
        print("✅ Text search working")

        ng.close()
        print("✅ Database operations test passed\n")
        return True

    except Exception as e:
        print(f"❌ Database operations test failed: {e}")
        return False
    finally:
        # Clean up
        if os.path.exists(db_path):
            os.unlink(db_path)

def test_embeddings_and_search():
    """Test embedding generation and semantic search."""
    print("🧪 Testing embeddings and semantic search...")

    with tempfile.NamedTemporaryFile(suffix='.db', delete=False) as tmp:
        db_path = tmp.name

    try:
        ng = NoteGraph(db_path)

        # Initialize OpenAI if available, otherwise use sentence transformers
        openai_key = os.environ.get("OPENAI_API_KEY")
        if openai_key:
            ng.init_openai(openai_key)
            ng.init_embedding_model(use_openai=True)
            print("✅ Using OpenAI embeddings")
        else:
            ng.init_embedding_model(use_openai=False)
            print("✅ Using sentence transformers (no API key)")

        # Add test notes
        ng.add_note(
            "Data Science",
            "Data science combines statistics, mathematics, and programming to extract insights from data.",
            ["Data Science", "Statistics", "Analytics"]
        )

        ng.add_note(
            "Artificial Intelligence",
            "AI refers to the simulation of human intelligence in machines that are programmed to think and learn.",
            ["AI", "Machine Learning", "Technology"]
        )

        ng.add_note(
            "Cooking Pasta",
            "To cook perfect pasta, bring water to boil, add salt, and cook for 8-10 minutes until al dente.",
            ["Cooking", "Food", "Recipes"]
        )

        # Test semantic search
        results = ng.semantic_search("machine learning algorithms", limit=3)
        assert len(results) > 0
        print(f"✅ Semantic search returned {len(results)} results")

        # Check if AI-related note appears first for ML query
        ml_results = ng.semantic_search("machine learning", limit=2)
        print(f"✅ ML search results: {[r['title'] for r in ml_results]}")

        ng.close()
        print("✅ Embeddings and search test passed\n")
        return True

    except Exception as e:
        print(f"❌ Embeddings and search test failed: {e}")
        return False
    finally:
        if os.path.exists(db_path):
            os.unlink(db_path)

def test_knowledge_graph():
    """Test knowledge graph functionality."""
    print("🧪 Testing knowledge graph functionality...")

    with tempfile.NamedTemporaryFile(suffix='.db', delete=False) as tmp:
        db_path = tmp.name

    try:
        ng = NoteGraph(db_path)
        ng.init_embedding_model(use_openai=True)

        # Add related notes
        ng.add_note(
            "Introduction to AI",
            "Artificial Intelligence is a field of computer science dedicated to creating systems that can perform tasks requiring human intelligence.",
            ["AI", "Introduction", "Technology"]
        )

        ng.add_note(
            "Deep Learning",
            "Deep learning is a subset of machine learning that uses neural networks with multiple layers to analyze various forms of data.",
            ["AI", "Deep Learning", "Neural Networks"]
        )

        # Test knowledge graph creation
        G = ng.get_knowledge_graph()
        assert len(G.nodes()) >= 2
        print(f"✅ Knowledge graph created with {len(G.nodes())} nodes and {len(G.edges())} edges")

        # Test note with connections
        conn = ng._get_connection()
        try:
            cursor = conn.execute("SELECT id FROM notes LIMIT 1")
            note_id = cursor.fetchone()[0]
        finally:
            conn.close()

        note_with_connections = ng.get_note_with_connections(note_id)
        assert 'connections' in note_with_connections
        print(f"✅ Retrieved note with connections: {len(note_with_connections['connections'])} connections")

        ng.close()
        print("✅ Knowledge graph test passed\n")
        return True

    except Exception as e:
        print(f"❌ Knowledge graph test failed: {e}")
        return False
    finally:
        if os.path.exists(db_path):
            os.unlink(db_path)

def test_llm_integration():
    """Test LLM integration for summarization and relationship generation."""
    print("🧪 Testing LLM integration...")

    with tempfile.NamedTemporaryFile(suffix='.db', delete=False) as tmp:
        db_path = tmp.name

    try:
        ng = NoteGraph(db_path)
        ng.init_embedding_model(use_openai=True)

        # Initialize OpenAI client
        openai_key = os.environ.get("OPENAI_API_KEY")
        if openai_key:
            ng.init_openai(openai_key)
            print("✅ OpenAI client initialized")
        else:
            print("⚠️  OpenAI API key not found, skipping OpenAI tests")
            # Use sentence transformers as fallback
            ng.init_embedding_model(use_openai=False)
            print("✅ Using sentence transformers as fallback")

        # Add notes for summarization
        note1_id = ng.add_note(
            "Climate Change Overview",
            "Climate change refers to long-term shifts in global temperatures and weather patterns. Human activities have been the main driver of climate change since the 1800s.",
            ["Climate", "Environment", "Science"]
        )

        note2_id = ng.add_note(
            "Renewable Energy Solutions",
            "Renewable energy sources like solar, wind, and hydroelectric power offer sustainable alternatives to fossil fuels and can help mitigate climate change.",
            ["Energy", "Sustainability", "Climate"]
        )

        # Test summarization (only if LLM is available)
        if ng.openai_client:
            summary = ng.summarize_notes([note1_id, note2_id])
            assert len(summary) > 50  # Summary should be substantial
            print("✅ LLM summarization working")
        else:
            print("⚠️  Skipping LLM summarization test (no API key)")

        ng.close()
        print("✅ LLM integration test passed\n")
        return True

    except Exception as e:
        print(f"❌ LLM integration test failed: {e}")
        return False
    finally:
        if os.path.exists(db_path):
            os.unlink(db_path)

def test_vector_operations():
    """Test specific vector operations from sqlite-vec."""
    print("🧪 Testing vector operations...")

    with tempfile.NamedTemporaryFile(suffix='.db', delete=False) as tmp:
        db_path = tmp.name

    try:
        import sqlite3
        import sqlite_vec

        # Test basic vector operations
        conn = sqlite3.connect(db_path)
        conn.enable_load_extension(True)
        sqlite_vec.load(conn)
        conn.enable_load_extension(False)

        # Test vec_version
        version = conn.execute("select vec_version()").fetchone()[0]
        print(f"✅ SQLite-Vec version: {version}")

        # Test vec_length
        test_vector = [0.1, 0.2, 0.3, 0.4]
        result = conn.execute('select vec_length(?)', [sqlite_vec.serialize_float32(test_vector)])
        length = result.fetchone()[0]
        assert length == 4
        print("✅ Vector length calculation working")

        # Test vector table creation and search
        conn.execute("create virtual table test_vec using vec0(embedding float[4])")
        conn.execute("insert into test_vec(rowid, embedding) values (1, '[0.1, 0.2, 0.3, 0.4]')")

        # Test vector search
        cursor = conn.execute("""
            select rowid, distance from test_vec
            where embedding match '[0.1, 0.2, 0.3, 0.4]'
            order by distance limit 1
        """)
        results = cursor.fetchall()
        assert len(results) > 0
        print("✅ Vector search operations working")

        conn.close()
        print("✅ Vector operations test passed\n")
        return True

    except Exception as e:
        print(f"❌ Vector operations test failed: {e}")
        return False
    finally:
        if os.path.exists(db_path):
            os.unlink(db_path)

def run_all_tests():
    """Run all tests and report results."""
    print("🚀 Starting NoteGraph Comprehensive Tests\n")

    tests = [
        test_vector_operations,
        test_database_operations,
        test_embeddings_and_search,
        test_knowledge_graph,
        test_llm_integration
    ]

    passed = 0
    total = len(tests)

    for test in tests:
        try:
            if test():
                passed += 1
        except Exception as e:
            print(f"❌ Test {test.__name__} failed with exception: {e}\n")

    print("="*50)
    print(f"🏆 Test Results: {passed}/{total} tests passed")

    if passed == total:
        print("🎉 All tests passed! NoteGraph is working correctly.")
        return True
    else:
        print(f"⚠️  {total - passed} tests failed. Please check the issues above.")
        return False

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