fix_embeddings.py

#!/usr/bin/env python3
"""
Script to regenerate embeddings for all notes after changing embedding model.
This fixes the "No results found" issue when switching to text-embedding-3-small.
"""

import sys
import os
from note_graph import NoteGraph

def main():
    print("🔧 Fixing embeddings after model change...")

    # Check if OpenAI API key is available
    if not os.getenv('OPENAI_API_KEY'):
        print("❌ OPENAI_API_KEY environment variable not found!")
        print("Please set your OpenAI API key:")
        print("export OPENAI_API_KEY='your-api-key-here'")
        return False

    try:
        # Initialize NoteGraph
        print("📚 Initializing NoteGraph...")
        ng = NoteGraph()

        # Check current state
        print("🔍 Checking current database state...")
        conn = ng.get_connection()
        cursor = conn.cursor()

        # Count notes
        cursor.execute("SELECT COUNT(*) FROM notes")
        notes_count = cursor.fetchone()[0]
        print(f"📝 Found {notes_count} notes in database")

        # Count current embeddings
        try:
            cursor.execute("SELECT COUNT(*) FROM note_embeddings")
            embeddings_count = cursor.fetchone()[0]
            print(f"🔢 Current embeddings: {embeddings_count}")
        except Exception as e:
            print(f"ℹ️  No embeddings table found: {e}")
            embeddings_count = 0

        if embeddings_count > 0:
            print("⚠️  Existing embeddings found. They will be recreated.")

        # Get all notes
        cursor.execute("SELECT id, title, content FROM notes")
        notes = cursor.fetchall()

        if not notes:
            print("❌ No notes found in database!")
            return False

        print(f"\n🔄 Regenerating embeddings for {len(notes)} notes...")

        # Clear existing embeddings
        cursor.execute("DELETE FROM note_embeddings")
        conn.commit()
        print("🧹 Cleared existing embeddings")

        # Regenerate embeddings for each note
        success_count = 0
        for i, (note_id, title, content) in enumerate(notes, 1):
            print(f"📊 Processing note {i}/{len(notes)}: {title[:50]}...")

            try:
                # Create full text for embedding
                full_text = f"{title}\n\n{content}"

                # Generate and store embedding
                ng._store_embedding(note_id, full_text)
                success_count += 1
                print(f"✅ Successfully generated embedding for note {note_id}")

            except Exception as e:
                print(f"❌ Failed to generate embedding for note {note_id}: {e}")

        print(f"\n🎉 Successfully generated {success_count}/{len(notes)} embeddings!")

        # Verify the embeddings
        cursor.execute("SELECT COUNT(*) FROM note_embeddings")
        final_count = cursor.fetchone()[0]
        print(f"✅ Final embedding count: {final_count}")

        # Test semantic search
        if final_count > 0:
            print("\n🔍 Testing semantic search...")
            try:
                # Try a simple search
                test_results = ng.semantic_search("AI", limit=3)
                print(f"🎯 Semantic search test: Found {len(test_results)} results")

                if test_results:
                    print("📋 Sample results:")
                    for i, result in enumerate(test_results[:3], 1):
                        print(f"  {i}. {result.get('title', 'Unknown')} (score: {result.get('score', 0):.3f})")

            except Exception as e:
                print(f"❌ Semantic search test failed: {e}")

        conn.close()
        return True

    except Exception as e:
        print(f"❌ Error fixing embeddings: {e}")
        return False

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