note_graph.py

import sqlite3
import sqlite_vec
import numpy as np
import json
import openai
import anthropic
from sentence_transformers import SentenceTransformer
import networkx as nx
import gradio as gr
import pandas as pd
from datetime import datetime
import os
from typing import List, Dict, Tuple, Optional

class NoteGraph:
    def __init__(self, db_path: str = "note_graph.db"):
        """Initialize the NoteGraph system with SQLite database and vector support."""
        self.db_path = db_path
        self.embedding_model = None
        self.openai_client = None
        self.anthropic_client = None
        self.use_openai_embeddings = False  # Default to False until initialized

        # Initialize database schema (connection created per operation)
        self._init_db_schema()

    def _init_db_schema(self):
        """Initialize database schema (creates tables if they don't exist)."""
        conn = self._get_connection()
        try:
            # Create notes table
            conn.execute("""
                CREATE TABLE IF NOT EXISTS notes (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    title TEXT NOT NULL,
                    content TEXT NOT NULL,
                    tags TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)

            # Check if vector table exists and handle dimension migration
            self._ensure_vector_table(conn)

            # Create knowledge graph relationships table
            conn.execute("""
                CREATE TABLE IF NOT EXISTS note_relationships (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    source_note_id INTEGER,
                    target_note_id INTEGER,
                    relationship_type TEXT,
                    strength REAL DEFAULT 1.0,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (source_note_id) REFERENCES notes (id),
                    FOREIGN KEY (target_note_id) REFERENCES notes (id)
                )
            """)

            conn.commit()
        finally:
            if conn:
                conn.close()

    def _ensure_vector_table(self, conn):
        """Ensure vector table exists with correct dimensions."""
        # Check if vector table exists
        result = conn.execute("""
            SELECT name FROM sqlite_master
            WHERE type='table' AND name='note_embeddings'
        """).fetchone()

        if result:
            # Table exists, check if we need to migrate
            try:
                # Test if we can access the table with current dimensions
                conn.execute("SELECT COUNT(*) FROM note_embeddings LIMIT 1").fetchone()
            except Exception as e:
                if "dimension mismatch" in str(e).lower() or "mismatch" in str(e).lower():
                    print("🔄 Migrating vector table to new dimensions...")
                    # Drop and recreate with correct dimensions
                    conn.execute("DROP TABLE note_embeddings")
                    self._create_vector_table(conn, self._get_embedding_dimensions())
                else:
                    raise e
        else:
            # Create new table with appropriate dimensions
            self._create_vector_table(conn, self._get_embedding_dimensions())

    def _create_vector_table(self, conn, dimensions):
        """Create vector table with specified dimensions."""
        conn.execute(f"""
            CREATE VIRTUAL TABLE note_embeddings
            USING vec0(note_embedding float[{dimensions}])
        """)
        print(f"✅ Created vector table with {dimensions} dimensions")

    def _get_embedding_dimensions(self):
        """Get the dimensions for the current embedding model."""
        # Check if we intend to use OpenAI embeddings, even if client isn't initialized yet
        if hasattr(self, 'use_openai_embeddings') and self.use_openai_embeddings:
            return 1536  # OpenAI text-embedding-3-small
        else:
            return 384   # Sentence transformers fallback

    def _get_connection(self):
        """Get a fresh database connection with vector support."""
        conn = sqlite3.connect(self.db_path)
        conn.enable_load_extension(True)
        sqlite_vec.load(conn)
        conn.enable_load_extension(False)
        return conn

    def init_embedding_model(self, use_openai: bool = True, model_name: str = "all-MiniLM-L6-v2"):
        """Initialize the embedding model for generating text embeddings."""
        self.use_openai_embeddings = use_openai
        if not use_openai:
            # Fallback to sentence transformers
            self.embedding_model = SentenceTransformer(model_name)
        else:
            # Initialize fallback model in case OpenAI fails
            self.embedding_model = SentenceTransformer(model_name)

    def init_openai(self, api_key: str):
        """Initialize OpenAI client."""
        self.openai_client = openai.OpenAI(api_key=api_key)

    def init_anthropic(self, api_key: str):
        """Initialize Anthropic client."""
        self.anthropic_client = anthropic.Anthropic(api_key=api_key)

    def add_note(self, title: str, content: str, tags: List[str] = None) -> int:
        """Add a new note to the database."""
        conn = self._get_connection()
        try:
            cursor = conn.execute(
                "INSERT INTO notes (title, content, tags) VALUES (?, ?, ?)",
                (title, content, json.dumps(tags or []))
            )
            note_id = cursor.lastrowid
            conn.commit()

            # Generate and store embedding
            self._generate_and_store_embedding(note_id, content)

            # Auto-generate relationships using LLM
            self._auto_generate_relationships(note_id, title, content)

            return note_id
        finally:
            if conn:
                conn.close()

    def _create_embedding_vector(self, text: str):
        """Create embedding vector for supplied text."""
        embedding = None

        if self.use_openai_embeddings and self.openai_client:
            try:
                response = self.openai_client.embeddings.create(
                    model="text-embedding-3-small",
                    input=text
                )
                embedding = response.data[0].embedding
                print(f"✅ Generated OpenAI embedding ({len(embedding)} dimensions)")
            except Exception as e:
                print(f"Error generating OpenAI embedding: {e}")
                if hasattr(self, 'embedding_model'):
                    embedding = self.embedding_model.encode(text)
                    print(f"✅ Fallback to sentence transformer ({len(embedding)} dimensions)")
                else:
                    print("❌ No embedding model available for fallback")
                    return None
        else:
            if hasattr(self, 'embedding_model'):
                embedding = self.embedding_model.encode(text)
                print(f"✅ Generated sentence transformer embedding ({len(embedding)} dimensions)")
            else:
                print("❌ No embedding model available")
                return None

        return embedding

    def _generate_and_store_embedding(self, note_id: int, content: str):
        """Generate embedding for note content and store in vector table."""
        embedding = self._create_embedding_vector(content)
        if embedding is None:
            print("❌ Failed to generate embedding")
            return

        # Store embedding with dimension mismatch handling
        try:
            embedding_bytes = sqlite_vec.serialize_float32(embedding)

            conn = self._get_connection()
            try:
                conn.execute(
                    "INSERT INTO note_embeddings (rowid, note_embedding) VALUES (?, ?)",
                    (note_id, embedding_bytes)
                )
                conn.commit()
                print(f"✅ Stored embedding for note {note_id}")
            except Exception as e:
                if "dimension mismatch" in str(e).lower() or "mismatch" in str(e).lower():
                    print(f"🔄 Dimension mismatch detected, recreating vector table...")
                    # Ensure the failed transaction is cleared before recreating the table
                    try:
                        conn.rollback()
                    except Exception:
                        pass
                    conn.close()
                    conn = None
                    # Recreate vector table with correct dimensions
                    self._recreate_vector_table_with_dimensions(len(embedding))
                    # Retry storing the embedding with a fresh connection
                    conn = self._get_connection()
                    conn.execute(
                        "INSERT INTO note_embeddings (rowid, note_embedding) VALUES (?, ?)",
                        (note_id, embedding_bytes)
                    )
                    conn.commit()
                    print(f"✅ Stored embedding after table recreation for note {note_id}")
                else:
                    raise e
            finally:
                if conn:
                    conn.close()
        except Exception as e:
            print(f"❌ Error storing embedding: {e}")

    def _recreate_vector_table_with_dimensions(self, dimensions):
        """Recreate vector table with specific dimensions."""
        conn = self._get_connection()
        try:
            # Drop existing table
            conn.execute("DROP TABLE IF EXISTS note_embeddings")
            conn.commit()

            # Create new table with correct dimensions
            conn.execute(f"""
                CREATE VIRTUAL TABLE note_embeddings
                USING vec0(note_embedding float[{dimensions}])
            """)
            conn.commit()
            print(f"✅ Recreated vector table with {dimensions} dimensions")
        except Exception as e:
            print(f"❌ Error recreating vector table: {e}")
            raise e
        finally:
            conn.close()

        # After recreation, repopulate embeddings so semantic search keeps working
        self._regenerate_all_embeddings()

    def _regenerate_all_embeddings(self):
        """Regenerate embeddings for every existing note (used after schema changes)."""
        conn = self._get_connection()
        try:
            rows = conn.execute("SELECT id, content FROM notes").fetchall()
        finally:
            conn.close()

        if not rows:
            print("ℹ️ No notes found while regenerating embeddings.")
            return

        print(f"🔁 Regenerating embeddings for {len(rows)} notes...")
        conn = self._get_connection()
        try:
            for note_id, content in rows:
                embedding = self._create_embedding_vector(content)
                if embedding is None:
                    print(f"⚠️ Skipping note {note_id}; embedding generation failed.")
                    continue
                embedding_bytes = sqlite_vec.serialize_float32(embedding)
                conn.execute(
                    "INSERT OR REPLACE INTO note_embeddings (rowid, note_embedding) VALUES (?, ?)",
                    (note_id, embedding_bytes)
                )
            conn.commit()
            print("✅ Finished regenerating embeddings.")
        except Exception as e:
            conn.rollback()
            print(f"❌ Error while regenerating embeddings: {e}")
            raise e
        finally:
            conn.close()

    def _auto_generate_relationships(self, note_id: int, title: str, content: str):
        """Auto-generate relationships with existing notes using LLM."""
        conn = self._get_connection()
        try:
            # Get existing notes
            existing_notes = conn.execute(
                "SELECT id, title, content FROM notes WHERE id != ? LIMIT 5",
                (note_id,)
            ).fetchall()

            if not existing_notes:
                return

            # Use LLM to find relationships
            relationships = self._find_relationships_with_llm(title, content, existing_notes)

            # Store relationships
            for rel in relationships:
                conn.execute(
                    "INSERT INTO note_relationships (source_note_id, target_note_id, relationship_type, strength) VALUES (?, ?, ?, ?)",
                    (note_id, rel['target_id'], rel['type'], rel['strength'])
                )
            conn.commit()
        finally:
            conn.close()

    def _find_relationships_with_llm(self, title: str, content: str, existing_notes: List[Tuple]) -> List[Dict]:
        """Use LLM to find relationships between new note and existing notes."""
        if not self.openai_client and not self.anthropic_client:
            return []

        # Prepare existing notes context
        notes_context = "\n".join([
            f"Note {nid}: {ntitle}\nContent: {ncontent[:200]}..."
            for nid, ntitle, ncontent in existing_notes
        ])

        prompt = f"""
        Analyze the following new note and find relationships with existing notes:

        New Note:
        Title: {title}
        Content: {content[:500]}...

        Existing Notes:
        {notes_context}

        Identify relationships and return as JSON array with format:
        [{{"target_id": note_id, "type": "relationship_type", "strength": 0.0-1.0}}]

        Relationship types: "similar", "related", "reference", "followup", "contrast"
        Only return relationships with strength > 0.6.
        """

        try:
            if self.openai_client:
                response = self.openai_client.chat.completions.create(
                    model="gpt-3.5-turbo",
                    messages=[{"role": "user", "content": prompt}],
                    temperature=0.3
                )
                result = response.choices[0].message.content
            else:
                response = self.anthropic_client.messages.create(
                    model="claude-3-haiku-20240307",
                    messages=[{"role": "user", "content": prompt}],
                    temperature=0.3,
                    max_tokens=1000
                )
                result = response.content[0].text

            # Parse JSON response
            relationships = json.loads(result)
            return relationships

        except Exception as e:
            print(f"Error generating relationships: {e}")
            return []

    def semantic_search(self, query: str, limit: int = 5) -> List[Dict]:
        """Perform semantic search on notes using vector similarity."""
        # Generate query embedding
        if self.use_openai_embeddings and self.openai_client:
            try:
                response = self.openai_client.embeddings.create(
                    model="text-embedding-3-small",
                    input=query
                )
                query_embedding = response.data[0].embedding
            except Exception as e:
                print(f"Error generating OpenAI embedding for query: {e}")
                # Fallback to sentence transformer if available
                if hasattr(self, 'embedding_model'):
                    query_embedding = self.embedding_model.encode(query)
                else:
                    return []
        else:
            # Use sentence transformer fallback
            if hasattr(self, 'embedding_model'):
                query_embedding = self.embedding_model.encode(query)
            else:
                return []

        query_bytes = sqlite_vec.serialize_float32(query_embedding)

        conn = self._get_connection()
        try:
            # Perform vector search with dimension mismatch handling
            try:
                cursor = conn.execute(f"""
                    SELECT
                        n.id, n.title, n.content, n.tags,
                        e.distance
                    FROM note_embeddings e
                    JOIN notes n ON e.rowid = n.id
                    WHERE note_embedding MATCH ?
                    AND k = {limit}
                """, (query_bytes,))
            except Exception as e:
                if "dimension mismatch" in str(e).lower() or "mismatch" in str(e).lower():
                    print(f"🔄 Search dimension mismatch, recreating vector table...")
                    # Release the current connection before dropping/recreating the table
                    try:
                        conn.rollback()
                    except Exception:
                        pass
                    conn.close()
                    conn = None
                    # Recreate vector table with query embedding dimensions
                    self._recreate_vector_table_with_dimensions(len(query_embedding))
                    # Retry the search with a new connection (results will be empty until embeddings are regenerated)
                    conn = self._get_connection()
                    cursor = conn.execute(f"""
                        SELECT
                            n.id, n.title, n.content, n.tags,
                            e.distance
                        FROM note_embeddings e
                        JOIN notes n ON e.rowid = n.id
                        WHERE note_embedding MATCH ?
                        AND k = {limit}
                    """, (query_bytes,))
                else:
                    raise e

            results = []

            for row in cursor.fetchall():
                # Convert cosine distance to similarity score (0-1 scale, higher is better)
                # distance is from 0 to 2, where 0 is perfect match
           
                print(f'{row[0]} {row[4]}')

                results.append({
                    'id': row[0],
                    'title': row[1],
                    'content': row[2],
                    'tags': json.loads(row[3]) if row[3] else [],
                    'similarity': row[4] 
                })
 

            return results
        finally:
            conn.close()

    def get_knowledge_graph(self) -> nx.Graph:
        """Build and return NetworkX graph from note relationships."""
        G = nx.Graph()
        conn = self._get_connection()
        try:
            # Add nodes (notes)
            notes = conn.execute("SELECT id, title FROM notes").fetchall()
            for note_id, title in notes:
                G.add_node(note_id, title=title)

            # Add edges (relationships)
            relationships = conn.execute("""
                SELECT source_note_id, target_note_id, relationship_type, strength
                FROM note_relationships
            """).fetchall()

            for source, target, rel_type, strength in relationships:
                G.add_edge(source, target, relationship=rel_type, weight=strength)

            return G
        finally:
            conn.close()

    def get_note_with_connections(self, note_id: int) -> Dict:
        """Get note details with its connections."""
        conn = self._get_connection()
        try:
            note = conn.execute(
                "SELECT * FROM notes WHERE id = ?", (note_id,)
            ).fetchone()

            if not note:
                return None

            # Get connections
            connections = conn.execute("""
                SELECT
                    nr.target_note_id,
                    n.title,
                    nr.relationship_type,
                    nr.strength
                FROM note_relationships nr
                JOIN notes n ON nr.target_note_id = n.id
                WHERE nr.source_note_id = ?
                ORDER BY nr.strength DESC
            """, (note_id,)).fetchall()

            return {
                'id': note[0],
                'title': note[1],
                'content': note[2],
                'tags': json.loads(note[3]) if note[3] else [],
                'created_at': note[4],
                'updated_at': note[5],
                'connections': [
                    {
                        'id': conn[0],
                        'title': conn[1],
                        'relationship': conn[2],
                        'strength': conn[3]
                    }
                    for conn in connections
                ]
            }
        finally:
            conn.close()

    def summarize_notes(self, note_ids: List[int]) -> str:
        """Generate a summary of multiple notes using LLM."""
        if not self.openai_client and not self.anthropic_client:
            return "LLM not initialized"

        conn = self._get_connection()
        try:
            # Get notes content
            notes_content = []
            for note_id in note_ids:
                note = conn.execute(
                    "SELECT title, content FROM notes WHERE id = ?", (note_id,)
                ).fetchone()
                if note:
                    notes_content.append(f"Title: {note[0]}\nContent: {note[1]}")

            if not notes_content:
                return "No notes found"

            prompt = f"""
            Please provide a comprehensive summary of these related notes:

            {'---'.join(notes_content)}

            Focus on:
            1. Key themes and concepts
            2. Important relationships
            3. Action items or insights
            4. Overall narrative or story
            """

            try:
                if self.openai_client:
                    response = self.openai_client.chat.completions.create(
                        model="gpt-3.5-turbo",
                        messages=[{"role": "user", "content": prompt}],
                        temperature=0.5,
                        max_tokens=800
                    )
                    return response.choices[0].message.content
                else:
                    response = self.anthropic_client.messages.create(
                        model="claude-3-haiku-20240307",
                        messages=[{"role": "user", "content": prompt}],
                        temperature=0.5,
                        max_tokens=800
                    )
                    return response.content[0].text

            except Exception as e:
                return f"Error generating summary: {e}"
        finally:
            conn.close()

    def remove_note(self, note_id: int) -> bool:
        """Remove a specific note and all its related data."""
        conn = self._get_connection()
        try:
            # Check if note exists
            note = conn.execute("SELECT id FROM notes WHERE id = ?", (note_id,)).fetchone()
            if not note:
                return False

            # Remove in order to handle foreign key constraints:
            # 1. Remove relationships where this note is source or target
            conn.execute("DELETE FROM note_relationships WHERE source_note_id = ? OR target_note_id = ?", (note_id, note_id))

            # 2. Remove embedding if it exists
            conn.execute("DELETE FROM note_embeddings WHERE rowid = ?", (note_id,))

            # 3. Remove the note itself
            conn.execute("DELETE FROM notes WHERE id = ?", (note_id,))

            conn.commit()
            return True
        except Exception as e:
            print(f"Error removing note {note_id}: {e}")
            conn.rollback()
            return False
        finally:
            conn.close()

    def remove_all_notes(self) -> int:
        """Remove all notes and reset the database."""
        conn = self._get_connection()
        try:
            # Get count before deletion
            count_result = conn.execute("SELECT COUNT(*) FROM notes").fetchone()
            total_notes = count_result[0] if count_result else 0

            # Delete all data in correct order
            conn.execute("DELETE FROM note_relationships")
            conn.execute("DELETE FROM note_embeddings")
            conn.execute("DELETE FROM notes")

            # Reset autoincrement
            conn.execute("DELETE FROM sqlite_sequence WHERE name = 'notes'")

            conn.commit()
            return total_notes
        except Exception as e:
            print(f"Error removing all notes: {e}")
            conn.rollback()
            return 0
        finally:
            conn.close()

    def get_all_note_ids(self) -> List[int]:
        """Get all note IDs for batch operations."""
        conn = self._get_connection()
        try:
            result = conn.execute("SELECT id FROM notes ORDER BY id").fetchall()
            return [row[0] for row in result]
        finally:
            conn.close()

    def chat_with_rag(self, message: str, conversation_history: List[Dict] = None, context_limit: int = 5) -> Dict:
        """Chat with RAG (Retrieval-Augmented Generation) using Graph Notes as context."""
        if conversation_history is None:
            conversation_history = []

        try:
            # Retrieve relevant notes using semantic search
            relevant_notes = self.semantic_search(message, limit=context_limit)

            # Build context from retrieved notes
            context = self._build_context_from_notes(relevant_notes)

            # Create system prompt for Ruby AI assistant
            system_prompt = """You are Ruby, an AI assistant specialized in helping users with their knowledge graph notes. You have access to the user's personal notes through a RAG (Retrieval-Augmented Generation) system.

Your role is to:
1. Help users understand and organize their notes
2. Answer questions based on their note content
3. Suggest connections between different notes
4. Assist with Ruby programming and general knowledge
5. Help users discover insights from their knowledge graph

When responding:
- Use the provided context from their notes to inform your answers
- If the context doesn't contain relevant information, say so clearly
- Be helpful, conversational, and engaging
- Reference specific notes when relevant (using note titles)
- Ask follow-up questions to help them explore their knowledge

Context from user's notes:
""" + context

            # Build conversation messages
            messages = [{"role": "system", "content": system_prompt}]

            # Add conversation history (limit to last 10 exchanges to manage context)
            for exchange in conversation_history[-10:]:
                if exchange.get("role") in ["user", "assistant"]:
                    messages.append({"role": exchange["role"], "content": exchange["content"]})

            # Add current user message
            messages.append({"role": "user", "content": message})

            # Generate response using available LLM
            response = self._generate_llm_response(messages)

            # Return response with metadata
            return {
                "response": response,
                "context_notes": relevant_notes,
                "context_used": len(relevant_notes) > 0
            }

        except Exception as e:
            return {
                "response": f"I apologize, but I encountered an error while processing your request: {str(e)}",
                "context_notes": [],
                "context_used": False
            }

    def _build_context_from_notes(self, notes: List[Dict]) -> str:
        """Build context string from retrieved notes."""
        if not notes:
            return "No relevant notes found."

        context_parts = []
        for i, note in enumerate(notes, 1):
            note_text = f"""
Note {i}: {note['title']}
Tags: {', '.join(note['tags']) if note['tags'] else 'None'}
Content: {note['content']}
Relevance: {note.get('similarity', 'N/A')}
"""
            context_parts.append(note_text)

        return "\n".join(context_parts)

    def _generate_llm_response(self, messages: List[Dict]) -> str:
        """Generate response using available LLM."""
        try:
            if self.openai_client:
                response = self.openai_client.chat.completions.create(
                    model="gpt-3.5-turbo",
                    messages=messages,
                    temperature=0.7,
                    max_tokens=1000
                )
                return response.choices[0].message.content
            elif self.anthropic_client:
                # Convert messages format for Anthropic
                anthropic_messages = []
                system_message = None

                for msg in messages:
                    if msg["role"] == "system":
                        system_message = msg["content"]
                    else:
                        anthropic_messages.append(msg)

                response = self.anthropic_client.messages.create(
                    model="claude-3-haiku-20240307",
                    messages=anthropic_messages,
                    system=system_message,
                    temperature=0.7,
                    max_tokens=1000
                )
                return response.content[0].text
            else:
                return "I apologize, but no AI service is currently configured. Please set up OpenAI or Anthropic API keys to use the chat functionality."

        except Exception as e:
            return f"I encountered an error generating my response: {str(e)}"

    def close(self):
        """Close database connection (no-op in thread-safe version)."""
        pass  # Connections are created per operation and closed automatically