ui.py

import gradio as gr
from models import LSP
from utils import logger, log_step
import asyncio
from typing import Optional

class LSPUI:
    def __init__(self):
        self.lsp: Optional[LSP] = None
        self.setup_interface()

    def validate_inputs(self, api_key: str, file: gr.File, instruction: str):
        """Validate input parameters"""
        if not api_key:
            return False, "Error: API key required"
        if not file:
            return False, "Error: No file uploaded"
        if not instruction:
            return False, "Error: No edit instruction"
        return True, None

    def read_file_content(self, file: gr.File) -> str:
        """Read content from file upload"""
        if isinstance(file, str):
            with open(file, 'r') as f:
                return f.read()
        return file.read().decode('utf-8')

    async def process_smart(self, api_key: str, file: gr.File, instruction: str):
        """Handle smart approach independently"""
        try:
            # Validate and initialize
            valid, error = self.validate_inputs(api_key, file, instruction)
            if not valid:
                return error, error, error
            
            if not self.lsp:
                self.lsp = LSP(api_key)
            
            content = self.read_file_content(file)
            log_step("UI", "Starting smart approach")
            
            # Process
            result, traces, elapsed_time = await self.lsp.edit_smart(content, instruction)
            log_step("UI", f"Smart approach completed in {elapsed_time:.2f}s")
            
            return (
                result,
                "\n".join(traces),
                f"✓ Completed in {elapsed_time:.2f}s"
            )
        except Exception as e:
            error_msg = f"Error: {str(e)}"
            log_step("UI", f"Error in smart approach: {error_msg}")
            return error_msg, error_msg, "✗ Failed"

    async def process_naive(self, api_key: str, file: gr.File, instruction: str):
        """Handle naive approach independently"""
        try:
            # Validate and initialize
            valid, error = self.validate_inputs(api_key, file, instruction)
            if not valid:
                return error, error, error
            
            if not self.lsp:
                self.lsp = LSP(api_key)
            
            content = self.read_file_content(file)
            log_step("UI", "Starting naive approach")
            
            # Process
            result, traces, elapsed_time = await self.lsp.edit_naive(content, instruction)
            log_step("UI", f"Naive approach completed in {elapsed_time:.2f}s")
            
            return (
                result,
                "\n".join(traces),
                f"✓ Completed in {elapsed_time:.2f}s"
            )
        except Exception as e:
            error_msg = f"Error: {str(e)}"
            log_step("UI", f"Error in naive approach: {error_msg}")
            return error_msg, error_msg, "✗ Failed"

    def setup_interface(self):
        """Setup the Gradio interface"""
        with gr.Blocks(title="LSP Comparison") as self.blocks:
            gr.Markdown("# LLM Selective Processing Demo")
            
            with gr.Row():
                api_key = gr.Textbox(label="OpenAI API Key", type="password")
            
            with gr.Row():
                file_upload = gr.File(
                    label="Upload Markdown File",
                    file_types=[".md", ".txt"]
                )
                instruction = gr.Textbox(label="What to edit")
            
            with gr.Row():
                smart_btn = gr.Button("Update with LSP")
                naive_btn = gr.Button("Update w/o LSP")
            
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### Smart Approach (Section-Aware)")
                    smart_result = gr.Textbox(label="Result", lines=10, value="")
                    smart_trace = gr.Textbox(label="Traces", lines=10, value="")
                    smart_status = gr.Textbox(label="Status", value="")
                
                with gr.Column():
                    gr.Markdown("### Naive Approach (Full Document)")
                    naive_result = gr.Textbox(label="Result", lines=10, value="")
                    naive_trace = gr.Textbox(label="Traces", lines=10, value="")
                    naive_status = gr.Textbox(label="Status", value="")

            # Set up independent event handlers
            smart_btn.click(
                fn=lambda: ("Working on it...", "Processing...", "⏳ Running..."),
                outputs=[smart_result, smart_trace, smart_status]
            ).then(
                fn=self.process_smart,
                inputs=[api_key, file_upload, instruction],
                outputs=[smart_result, smart_trace, smart_status]
            )

            naive_btn.click(
                fn=lambda: ("Working on it...", "Processing...", "⏳ Running..."),
                outputs=[naive_result, naive_trace, naive_status]
            ).then(
                fn=self.process_naive,
                inputs=[api_key, file_upload, instruction],
                outputs=[naive_result, naive_trace, naive_status]
            )

    def launch(self):
        """Launch the Gradio interface"""
        self.blocks.launch()