gradio_app.py

import gradio as gr
import plotly.graph_objects as go
import plotly.express as px
import networkx as nx
import json
from note_graph import NoteGraph
import os
from typing import List, Tuple

class NoteGraphUI:
    def __init__(self):
        """Initialize the Gradio interface for NoteGraph."""
        self.ng = NoteGraph()
        self.conversation_history = []

        # Initialize LLM clients if API keys are available
        openai_key = os.getenv("OPENAI_API_KEY")
        anthropic_key = os.getenv("ANTHROPIC_API_KEY")

        if openai_key:
            self.ng.init_openai(openai_key)
            # Use OpenAI embeddings by default when API key is available
            self.ng.init_embedding_model(use_openai=True)
        else:
            # Fallback to sentence transformers if no API key
            self.ng.init_embedding_model(use_openai=False)

        if anthropic_key:
            self.ng.init_anthropic(anthropic_key)

    def add_note_ui(self, title: str, content: str, tags: str) -> str:
        """Add a new note via UI."""
        if not title or not content:
            return "❌ Title and content are required"

        try:
            tag_list = [tag.strip() for tag in tags.split(",")] if tags else []
            note_id = self.ng.add_note(title, content, tag_list)
            return f"✅ Note added successfully! ID: {note_id}"
        except Exception as e:
            return f"❌ Error adding note: {str(e)}"

    def search_notes_ui(self, query: str, search_type: str) -> List[List[str]]:
        """Search notes via UI."""
        if not query:
            return [["No results", "No content", "No tags", "0.0"]]

        try:
            if search_type == "Semantic Search":
                results = self.ng.semantic_search(query)
                if results:
                    return [
                        [
                            f"📝 {r['title']} (ID: {r['id']})",
                            r['content'][:200] + "..." if len(r['content']) > 200 else r['content'],
                            ", ".join(r['tags']) if r['tags'] else "No tags",
                            f"Similarity: {r['similarity']:.3f}"
                        ]
                        for r in results
                    ]
            else:
                # Text search - use a new connection
                conn = self.ng._get_connection()
                try:
                    cursor = conn.execute("""
                        SELECT id, title, content, tags FROM notes
                        WHERE title LIKE ? OR content LIKE ?
                        ORDER BY id DESC
                        LIMIT 10
                    """, (f"%{query}%", f"%{query}%"))

                    results = cursor.fetchall()
                    if results:
                        return [
                            [
                                f"📝 {row[1]} (ID: {row[0]})",
                                row[2][:200] + "..." if len(row[2]) > 200 else row[2],
                                ", ".join(json.loads(row[3])) if row[3] else "No tags",
                                "Text match"
                            ]
                            for row in results
                        ]
                finally:
                    conn.close()

            return [["No results found", "", "", ""]]
        except Exception as e:
            return [[f"❌ Search error: {str(e)}", "", "", ""]]

    def get_all_notes_ui(self) -> List[List[str]]:
        """Get all notes for display."""
        conn = self.ng._get_connection()
        try:
            cursor = conn.execute("""
                SELECT id, title, created_at FROM notes
                ORDER BY updated_at DESC
                LIMIT 20
            """)

            notes = cursor.fetchall()
            return [
                [f"📝 {note[1]} (ID: {note[0]})", note[2][:19]]
                for note in notes
            ]
        except Exception as e:
            return [[f"❌ Error: {str(e)}", ""]]
        finally:
            conn.close()

    def visualize_graph_ui(self) -> go.Figure:
        """Create interactive graph visualization."""
        try:
            G = self.ng.get_knowledge_graph()
            print(f"Debug: Graph has {len(G.nodes())} nodes and {len(G.edges())} edges")

            if len(G.nodes()) == 0:
                # Return empty plot
                fig = go.Figure()
                fig.add_annotation(
                    text="No notes to visualize. Add some notes first!",
                    xref="paper", yref="paper",
                    x=0.5, y=0.5, showarrow=False,
                    font=dict(size=16)
                )
                return fig

            # Calculate layout with error handling
            try:
                pos = nx.spring_layout(G, k=1, iterations=50)
            except Exception as layout_error:
                print(f"Debug: Layout error: {layout_error}")
                # Fallback to circular layout if spring layout fails
                try:
                    pos = nx.circular_layout(G)
                except:
                    # Final fallback - manual positions
                    nodes = list(G.nodes())
                    pos = {node: (i % 5, i // 5) for i, node in enumerate(nodes)}

            # Extract node and edge positions with error handling
            node_x = []
            node_y = []
            node_text = []
            node_info = []

            for node in G.nodes():
                try:
                    x, y = pos[node]
                    node_x.append(float(x))
                    node_y.append(float(y))
                    title = G.nodes[node].get('title', f'Note {node}')
                    node_text.append(str(title)[:20] + '...' if len(str(title)) > 20 else str(title))
                    node_info.append(f"ID: {node}<br>Title: {title}")
                except Exception as node_error:
                    print(f"Debug: Error processing node {node}: {node_error}")
                    continue

            # Create edge traces with error handling
            edge_x = []
            edge_y = []

            for edge in G.edges():
                try:
                    x0, y0 = pos[edge[0]]
                    x1, y1 = pos[edge[1]]
                    edge_x.extend([float(x0), float(x1), None])
                    edge_y.extend([float(y0), float(y1), None])
                except Exception as edge_error:
                    print(f"Debug: Error processing edge {edge}: {edge_error}")
                    continue

            # Create figure
            fig = go.Figure()

            # Add edges (only if we have edges)
            if edge_x and edge_y:
                fig.add_trace(go.Scatter(
                    x=edge_x, y=edge_y,
                    line=dict(width=1, color='#888'),
                    hoverinfo='none',
                    mode='lines'
                ))

            # Add nodes (only if we have nodes)
            if node_x and node_y:
                fig.add_trace(go.Scatter(
                    x=node_x, y=node_y,
                    mode='markers+text',
                    hoverinfo='text',
                    text=node_text,
                    hovertext=node_info,
                    marker=dict(
                        size=20,
                        color='lightblue',
                        line=dict(width=2, color='darkblue')
                    ),
                    textposition="middle center"
                ))

            # Final check - if we have no nodes after processing, show appropriate message
            if not node_x:
                fig = go.Figure()
                fig.add_annotation(
                    text="No valid notes to display. Try adding some notes first.",
                    xref="paper", yref="paper",
                    x=0.5, y=0.5, showarrow=False,
                    font=dict(size=16)
                )
                return fig

            fig.update_layout(
                title=dict(
                    text="Knowledge Graph Visualization",
                    font=dict(size=18)
                ),
                showlegend=False,
                hovermode='closest',
                margin=dict(b=20, l=5, r=5, t=50),
                annotations=[
                    dict(
                        text="Interactive knowledge graph - hover over nodes for details",
                        showarrow=False,
                        xref="paper",
                        yref="paper",
                        x=0.005,
                        y=-0.002,
                        xanchor='left',
                        yanchor='bottom',
                        font=dict(color='#888', size=12)
                    )
                ],
                xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
                yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
                height=600
            )

            return fig

        except Exception as e:
            print(f"Debug: Knowledge graph visualization error: {e}")
            import traceback
            traceback.print_exc()
            # Return error plot
            fig = go.Figure()
            fig.add_annotation(
                text=f"Error creating visualization: {str(e)}<br><br>Please try refreshing the graph or check the console for details.",
                xref="paper", yref="paper",
                x=0.5, y=0.5, showarrow=False,
                font=dict(size=14, color="red")
            )
            return fig

    def get_note_details_ui(self, note_id: str) -> Tuple[str, str]:
        """Get detailed note information."""
        if not note_id:
            return "❌ Please enter a note ID", ""

        try:
            note_id = int(note_id)
            note = self.ng.get_note_with_connections(note_id)

            if not note:
                return "❌ Note not found", ""

            # Format note details
            details = f"""📝 **{note['title']}**
**ID:** {note['id']}
**Tags:** {', '.join(note['tags']) if note['tags'] else 'None'}
**Created:** {note['created_at'][:19]}
**Updated:** {note['updated_at'][:19]}

**Content:**
{note['content']}

**Connections ({len(note['connections'])}):**"""

            connections = ""
            for conn in note['connections']:
                connections += f"• {conn['title']} (ID: {conn['id']}) - {conn['relationship']} (strength: {conn['strength']:.2f})\n"

            return details, connections if connections else "No connections found"

        except ValueError:
            return "❌ Invalid note ID. Please enter a number.", ""
        except Exception as e:
            return f"❌ Error: {str(e)}", ""

    def generate_summary_ui(self, note_ids: str) -> str:
        """Generate summary of selected notes."""
        if not note_ids:
            return "❌ Please enter note IDs (comma-separated)"

        try:
            id_list = [int(id.strip()) for id in note_ids.split(",")]
            summary = self.ng.summarize_notes(id_list)
            return summary
        except ValueError:
            return "❌ Invalid note IDs. Please enter comma-separated numbers."
        except Exception as e:
            return f"❌ Error generating summary: {str(e)}"

    def remove_note_ui(self, note_id: str) -> str:
        """Remove a specific note via UI."""
        if not note_id:
            return "❌ Please enter a note ID"

        try:
            note_id = int(note_id.strip())

            # Check if note exists first
            note = self.ng.get_note_with_connections(note_id)
            if not note:
                return f"❌ Note with ID {note_id} not found"

            # Show confirmation details
            title = note['title']

            # Remove the note
            success = self.ng.remove_note(note_id)

            if success:
                return f"✅ Successfully removed note: '{title}' (ID: {note_id})"
            else:
                return f"❌ Failed to remove note with ID {note_id}"

        except ValueError:
            return "❌ Invalid note ID. Please enter a number."
        except Exception as e:
            return f"❌ Error removing note: {str(e)}"

    def remove_all_notes_ui(self, confirm: str) -> str:
        """Remove all notes via UI with confirmation."""
        if confirm.lower() != "delete all":
            return "❌ Please type 'delete all' exactly to confirm removal of all notes"

        try:
            # Get current note count
            note_ids = self.ng.get_all_note_ids()
            count = len(note_ids)

            if count == 0:
                return "ℹ️ No notes to remove"

            # Remove all notes
            removed_count = self.ng.remove_all_notes()

            return f"✅ Successfully removed {removed_count} notes and all their relationships, embeddings, and connections"

        except Exception as e:
            return f"❌ Error removing all notes: {str(e)}"

    def chat_with_ruby_ui(self, message: str, history: List, context_limit: int) -> Tuple[List, str, str]:
        """Chat with Ruby AI using RAG system."""
        if not message.strip():
            return history, "Please enter a message.", ""

        try:
            # Convert Gradio history format to our format
            conversation_history = []
            for user_msg, assistant_msg in history:
                conversation_history.append({"role": "user", "content": user_msg})
                if assistant_msg:
                    conversation_history.append({"role": "assistant", "content": assistant_msg})

            # Get response from Ruby AI
            result = self.ng.chat_with_rag(message, conversation_history, context_limit)

            # Add user message and response to history
            history.append([message, result["response"]])

            # Format context information
            context_info = ""
            if result["context_used"] and result["context_notes"]:
                context_info = f"📚 Used {len(result['context_notes'])} notes as context:\n"
                for note in result["context_notes"][:3]:  # Show top 3 notes
                    context_info += f"• {note['title']} (similarity: {note.get('similarity', 'N/A')})\n"
            else:
                context_info = "💭 No specific notes were used as context for this response."

            return history, "", context_info

        except Exception as e:
            error_msg = f"Error: {str(e)}"
            history.append([message, f"I apologize, but I encountered an error: {error_msg}"])
            return history, "", "❌ Error occurred during chat."

    def clear_chat_ui(self) -> Tuple[List, str]:
        """Clear chat history."""
        self.conversation_history = []
        return [], "Chat history cleared."

    def get_chat_status_ui(self) -> str:
        """Get chat system status."""
        openai_key = os.getenv("OPENAI_API_KEY")
        anthropic_key = os.getenv("ANTHROPIC_API_KEY")

        if openai_key or anthropic_key:
            provider = "OpenAI" if openai_key else "Anthropic"
            return f"✅ Ruby AI is ready! Using {provider} for responses.\n💬 You can ask questions about your notes, request summaries, or get help with Ruby programming."
        else:
            return "⚠️ Ruby AI requires an API key to function.\nPlease set OPENAI_API_KEY or ANTHROPIC_API_KEY environment variable to enable chat functionality."

    def create_interface(self):
        """Create the Gradio interface."""
        with gr.Blocks(title="NoteGraph - Knowledge Graph Notes") as interface:
            gr.Markdown("# 🧠 NoteGraph - Knowledge Graph Note Taking System")
            gr.Markdown("A semantic note-taking app with knowledge graph capabilities and vector search")

            with gr.Tabs():
                # Tab 1: Add Notes
                with gr.Tab("📝 Add Notes"):
                    with gr.Row():
                        with gr.Column():
                            title_input = gr.Textbox(label="Title", placeholder="Enter note title...")
                            content_input = gr.TextArea(label="Content", placeholder="Enter note content...", lines=10)
                            tags_input = gr.Textbox(label="Tags (comma-separated)", placeholder="tag1, tag2, tag3...")
                            add_btn = gr.Button("Add Note", variant="primary")
                            add_output = gr.Textbox(label="Status", interactive=False)

                        with gr.Column():
                            gr.Markdown("### 📊 Recent Notes")
                            notes_df = gr.DataFrame(
                                headers=["Note", "Created"],
                                datatype=["str", "str"],
                                interactive=False
                            )
                            refresh_notes_btn = gr.Button("Refresh Notes")

                # Tab 2: Search
                with gr.Tab("🔍 Search Notes"):
                    search_query = gr.Textbox(label="Search Query", placeholder="Enter search terms...")
                    search_type = gr.Radio(["Semantic Search", "Text Search"], label="Search Type", value="Semantic Search")
                    search_btn = gr.Button("Search")
                    search_results = gr.DataFrame(
                        headers=["Title", "Content Preview", "Tags", "Score"],
                        datatype=["str", "str", "str", "str"],
                        interactive=False
                    )

                # Tab 3: Knowledge Graph
                with gr.Tab("🕸️ Knowledge Graph"):
                    gr.Markdown("### Interactive Knowledge Graph Visualization")
                    graph_plot = gr.Plot(label="Knowledge Graph")
                    refresh_graph_btn = gr.Button("Refresh Graph")

                # Tab 4: Note Details & Analysis
                with gr.Tab("📊 Note Details & Analysis"):
                    with gr.Row():
                        with gr.Column():
                            note_id_input = gr.Textbox(label="Note ID", placeholder="Enter note ID...")
                            get_details_btn = gr.Button("Get Details")
                            note_details = gr.Textbox(label="Note Details", lines=15, interactive=False)
                            note_connections = gr.Textbox(label="Connections", lines=10, interactive=False)

                        with gr.Column():
                            summary_ids = gr.Textbox(label="Note IDs for Summary (comma-separated)",
                                                   placeholder="1, 2, 3...")
                            generate_summary_btn = gr.Button("Generate Summary")
                            summary_output = gr.TextArea(label="Generated Summary", lines=15, interactive=False)

                # Tab 5: Ruby AI Chat
                with gr.Tab("🤖 Ruby AI Chat"):
                    gr.Markdown("### 💬 Chat with Ruby AI - Your Knowledge Assistant")
                    gr.Markdown("Ruby AI uses your Graph Notes as context to provide personalized assistance. Ask questions about your notes, request summaries, or get help with Ruby programming!")

                    with gr.Row():
                        with gr.Column(scale=3):
                            # Chat interface
                            chatbot = gr.Chatbot(
                                label="Ruby AI Conversation",
                                height=500,
                                show_copy_button=True,
                                bubble_full_width=False
                            )

                            with gr.Row():
                                msg_input = gr.Textbox(
                                    label="Your Message",
                                    placeholder="Ask Ruby anything about your notes or Ruby programming...",
                                    lines=2,
                                    scale=4
                                )
                                send_btn = gr.Button("Send", variant="primary", scale=1)

                            with gr.Row():
                                clear_btn = gr.Button("🗑️ Clear Chat", variant="secondary")
                                context_limit = gr.Slider(
                                    minimum=1,
                                    maximum=10,
                                    value=5,
                                    step=1,
                                    label="Context Notes Limit",
                                    info="How many notes to use as context"
                                )

                        with gr.Column(scale=1):
                            # Status and context panel
                            gr.Markdown("### 📊 Chat Status")
                            chat_status = gr.Textbox(
                                label="System Status",
                                value=self.get_chat_status_ui(),
                                lines=3,
                                interactive=False
                            )

                            gr.Markdown("### 📚 Context Used")
                            context_display = gr.Textbox(
                                label="Retrieved Context",
                                value="Start a conversation to see context information...",
                                lines=8,
                                interactive=False,
                                max_lines=12
                            )

                            gr.Markdown("### 💡 Tips")
                            gr.Markdown("""
                            • Ask specific questions about your notes
                            • Request summaries of topics
                            • Explore connections between ideas
                            • Get Ruby programming help
                            • Use natural language - Ruby understands context!
                            """)

                # Tab 6: Remove Knowledge
                with gr.Tab("🗑️ Remove Knowledge"):
                    gr.Markdown("### Remove Notes and Knowledge")
                    gr.Markdown("⚠️ **Warning**: These actions cannot be undone!")

                    with gr.Row():
                        with gr.Column():
                            gr.Markdown("#### Remove Single Note")
                            remove_note_id = gr.Textbox(
                                label="Note ID to Remove",
                                placeholder="Enter the ID of the note to remove..."
                            )
                            remove_note_btn = gr.Button("Remove Note", variant="stop")
                            remove_note_output = gr.Textbox(
                                label="Removal Status",
                                interactive=False,
                                lines=2
                            )

                        with gr.Column():
                            gr.Markdown("#### Remove All Notes")
                            gr.Markdown("**DANGER ZONE**: This will delete all notes, embeddings, and relationships!")
                            remove_all_confirm = gr.Textbox(
                                label="Type 'delete all' to confirm",
                                placeholder="Type 'delete all' exactly..."
                            )
                            remove_all_btn = gr.Button("🚨 DELETE ALL NOTES", variant="stop", size="lg")
                            remove_all_output = gr.Textbox(
                                label="Mass Deletion Status",
                                interactive=False,
                                lines=3
                            )

                    # Refresh other components after removal
                    gr.Markdown("#### Refresh Other Tabs")
                    gr.Markdown("After removing notes, click the refresh buttons in other tabs to update the display.")
                    refresh_all_btn = gr.Button("🔄 Refresh All Components", variant="secondary")

            # Event handlers
            add_btn.click(
                fn=self.add_note_ui,
                inputs=[title_input, content_input, tags_input],
                outputs=add_output
            )

            refresh_notes_btn.click(
                fn=self.get_all_notes_ui,
                outputs=notes_df
            )

            search_btn.click(
                fn=self.search_notes_ui,
                inputs=[search_query, search_type],
                outputs=search_results
            )

            refresh_graph_btn.click(
                fn=self.visualize_graph_ui,
                outputs=graph_plot
            )

            get_details_btn.click(
                fn=self.get_note_details_ui,
                inputs=note_id_input,
                outputs=[note_details, note_connections]
            )

            generate_summary_btn.click(
                fn=self.generate_summary_ui,
                inputs=summary_ids,
                outputs=summary_output
            )

            # Remove functionality event handlers
            remove_note_btn.click(
                fn=self.remove_note_ui,
                inputs=remove_note_id,
                outputs=remove_note_output
            )

            remove_all_btn.click(
                fn=self.remove_all_notes_ui,
                inputs=remove_all_confirm,
                outputs=remove_all_output
            )

            refresh_all_btn.click(
                fn=self.get_all_notes_ui,
                outputs=notes_df
            )

            # Also refresh graph when refresh all is clicked
            refresh_all_btn.click(
                fn=self.visualize_graph_ui,
                outputs=graph_plot
            )

            # Ruby AI Chat event handlers
            msg_input.submit(
                fn=self.chat_with_ruby_ui,
                inputs=[msg_input, chatbot, context_limit],
                outputs=[chatbot, msg_input, context_display]
            )

            send_btn.click(
                fn=self.chat_with_ruby_ui,
                inputs=[msg_input, chatbot, context_limit],
                outputs=[chatbot, msg_input, context_display]
            )

            clear_btn.click(
                fn=self.clear_chat_ui,
                outputs=[chatbot, chat_status]
            )

            # Load initial data
            interface.load(
                fn=self.get_all_notes_ui,
                outputs=notes_df
            )

            interface.load(
                fn=self.visualize_graph_ui,
                outputs=graph_plot
            )

        return interface

    def launch(self, **kwargs):
        """Launch the Gradio interface."""
        interface = self.create_interface()
        interface.launch(**kwargs)

if __name__ == "__main__":
    # Set API keys (you can also set these as environment variables)
    app = NoteGraphUI()
    app.launch(server_name="0.0.0.0", server_port=7860, share=True)