chore: use imports to reduce app length

app.py

@@ -2,11 +2,6 @@ import os
 import sys
 import gradio as gr
 from dotenv import load_dotenv
-import subprocess
-import traceback
-import time
-import pickle
-import faiss
 
 # Add the current directory to the path so we can import from src
 sys.path.append(".")
@@ -22,1209 +17,57 @@ if not api_key:
 else:
     print("OpenAI API key loaded successfully.")
 
-# Global resource manager to persist data across requests
-class ResourceManager:
-    _instance = None
-    
-    def __new__(cls):
-        if cls._instance is None:
-            cls._instance = super(ResourceManager, cls).__new__(cls)
-            cls._instance.faiss_index = None
-            cls._instance.doc_chunks = None
-            cls._instance.embedding_vectors = None
-            cls._instance.initialized = False
-        return cls._instance
-    
-    def load_resources(self):
-        """Load all required resources"""
-        if self.initialized:
-            print("Resources already loaded, skipping...")
-            return True
-            
-        success = True
-        # Load FAISS index
-        try:
-            print("Loading FAISS index...")
-            self.faiss_index = faiss.read_index("embeddings/faiss_index.index")
-            print("FAISS index loaded successfully")
-        except Exception as e:
-            print(f"Error loading FAISS index: {e}")
-            success = False
-        
-        # Load document chunks
-        try:
-            print("Loading document chunks...")
-            with open("data/doc_chunks.pkl", "rb") as f:
-                self.doc_chunks = pickle.load(f)
-            print(f"Loaded {len(self.doc_chunks)} document chunks")
-        except Exception as e:
-            print(f"Error loading document chunks: {e}")
-            success = False
-            
-        # Load embeddings if available
-        try:
-            print("Loading embeddings...")
-            with open("embeddings/embeddings.pkl", "rb") as f:
-                self.embedding_vectors = pickle.load(f)
-            print("Embeddings loaded successfully")
-        except Exception as e:
-            print(f"Error loading embeddings: {e}")
-            # This is not critical, so don't set success to False
-            
-        if success:
-            self.initialized = True
-            print("All critical resources loaded successfully!")
-        
-        return success
-    
-    def get_faiss_index(self):
-        """Get the FAISS index, loading if necessary"""
-        if not self.initialized:
-            self.load_resources()
-        return self.faiss_index
-    
-    def get_doc_chunks(self):
-        """Get the document chunks, loading if necessary"""
-        if not self.initialized:
-            self.load_resources()
-        return self.doc_chunks
-    
-    def get_embedding_vectors(self):
-        """Get the embedding vectors, loading if necessary"""
-        if not self.initialized:
-            self.load_resources()
-        return self.embedding_vectors
+# Import our custom modules
+from src.utils.styles import get_custom_css
+from src.utils.ui import create_interface, create_fallback_ui
+from src.utils.resource_manager import ResourceManager, check_data_files
+from src.utils.query_handler import process_query_with_status
+
+# Check if data files exist and load them
+data_ready = check_data_files()
 
 # Create the global resource manager
 resource_manager = ResourceManager()
 
-# Before initializing your model/agent, check and download data
-data_files = [
-    "embeddings/faiss_index.index", 
-    "data/doc_chunks.pkl", 
-    "embeddings/embeddings.pkl"
-]
-
-missing_files = [f for f in data_files if not os.path.exists(f)]
-data_ready = True
-
-if missing_files:
-    print(f"Missing data files: {missing_files}")
-    print("Downloading or creating required data files...")
-    try:
-        result = subprocess.run(
-            [sys.executable, "download_from_hub.py"],
-            check=False,
-            capture_output=True,
-            text=True
-        )
-        print(result.stdout)
-        if result.returncode != 0:
-            print(f"Warning: Data preparation finished with return code {result.returncode}")
-            print(f"Error output: {result.stderr}")
-            data_ready = False
-    except Exception as e:
-        print(f"Error preparing data: {e}")
-        traceback.print_exc()
-        data_ready = False
-
-# Verify files exist before importing modules that need them
-if not all(os.path.exists(f) for f in data_files):
-    print("Warning: Some required data files are still missing.")
-    print("The application may not function correctly.")
-    data_ready = False
-else:
-    # Load resources into the resource manager
-    data_ready = resource_manager.load_resources()
-
-# Define a function to create fallback UI
-def create_fallback_ui():
-    with gr.Blocks(
-        theme=gr.themes.Base(),
-        css=custom_css,
-        title="Defensor Legal Assistant - Data Error"
-    ) as demo:
-        # Create a header with app-header class for better styling
-        with gr.Column(elem_classes="app-header"):
-            gr.HTML("""
-            <div class="logo-container">
-                <div class="logo">
-                    <span class="logo-icon">⚖️</span>
-                    <span class="logo-text">Defensor</span>
-                </div>
-            </div>
-            """)
-            gr.Markdown("# Agentic Defensor - Data Loading Error")
-        
-        with gr.Column(elem_classes="main-container"):
-            gr.Markdown("""
-            ### There was an error loading the required data files.
-            
-            This application requires specific dataset files to function properly. 
-            The system tried to download them from Hugging Face Hub but encountered errors.
-            
-            Please check the application logs for more details on what went wrong.
-            
-            Common issues:
-            1. Dataset repository might not be accessible or might not contain the expected files
-            2. Network connectivity issues when downloading datasets
-            3. Missing or invalid Hugging Face token for accessing private datasets
-            """)
-            
-            with gr.Row():
-                input_text = gr.Textbox(label="You can enter text here, but functionality is limited")
-                output_text = gr.Textbox(label="System response")
-                
-            submit_btn = gr.Button("Submit (Limited functionality)")
-            
-            def dummy_response(text):
-                return "The system is in fallback mode due to data loading errors. Please check application logs."
-                
-            submit_btn.click(dummy_response, inputs=[input_text], outputs=[output_text])
-        
-    return demo
-
 # Try importing modules
 director = None
 CHAT_MODEL = None
 try:
-    # Use a workaround to avoid circular imports - don't import here
     # We'll import these in the specific functions where needed
     data_ready = data_ready and os.path.exists('src/agents/agent_director.py')
 except ImportError as e:
     print(f"Error importing modules: {e}")
     print("This may be caused by missing data files or module dependencies.")
     data_ready = False
-    traceback.print_exc()
-
-# Define custom CSS for better UI
-custom_css = """
-body {
-    background-color: #f8f9fa;
-    color: #333;
-    font-family: 'Inter', -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif;
-}
-
-/* App-wide typography improvements */
-h1, h2, h3, h4, h5, h6 {
-    font-family: 'Inter', -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif;
-    font-weight: 600;
-    letter-spacing: -0.02em;
-}
-
-/* Dark mode theme - beautiful gradient background */
-.gradio-container {
-    background: linear-gradient(135deg, #172a46 0%, #0c1220 100%);
-    max-width: 1000px;
-    margin: 0 auto;
-    border-radius: 16px;
-    overflow: hidden;
-    box-shadow: 0 10px 30px rgba(0, 0, 0, 0.3);
-    border: 1px solid rgba(255, 255, 255, 0.05);
-}
-
-/* Logo styling */
-.logo-container {
-    text-align: center;
-    margin-bottom: 10px;
-}
-
-.logo {
-    display: inline-flex;
-    align-items: center;
-    background: rgba(0, 0, 0, 0.2);
-    padding: 10px 20px;
-    border-radius: 50px;
-    border: 1px solid rgba(255, 255, 255, 0.1);
-    box-shadow: 0 4px 15px rgba(0, 0, 0, 0.2);
-}
-
-.logo-icon {
-    font-size: 2rem;
-    margin-right: 10px;
-    display: inline-block;
-}
-
-.logo-text {
-    font-size: 1.6rem;
-    font-weight: 700;
-    color: white;
-    letter-spacing: 0.5px;
-    background: linear-gradient(90deg, #a5b4fc 0%, #818cf8 100%);
-    -webkit-background-clip: text;
-    -webkit-text-fill-color: transparent;
-    text-shadow: 0 2px 4px rgba(0, 0, 0, 0.2);
-}
-
-/* App header styling */
-.app-header {
-    text-align: center;
-    padding: 30px 20px 20px;
-    background: rgba(0, 0, 0, 0.2);
-    border-bottom: 1px solid rgba(255, 255, 255, 0.1);
-    margin-bottom: 15px;
-}
-
-.app-header h1 {
-    font-size: 2.4rem;
-    color: #fff;
-    margin-bottom: 10px;
-    font-weight: 700;
-    text-shadow: 0 2px 4px rgba(0, 0, 0, 0.3);
-    background: linear-gradient(90deg, #e0e7ff 0%, #c7d2fe 100%);
-    -webkit-background-clip: text;
-    -webkit-text-fill-color: transparent;
-}
-
-.app-header p {
-    color: rgba(255, 255, 255, 0.8);
-    font-size: 1.1rem;
-    max-width: 700px;
-    margin: 0 auto;
-    line-height: 1.5;
-}
-
-/* Main container styling */
-.main-container {
-    background-color: rgba(30, 41, 59, 0.7);
-    backdrop-filter: blur(10px);
-    border-radius: 12px;
-    padding: 25px;
-    margin: 15px;
-    color: #fff;
-    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.15);
-    border: 1px solid rgba(255, 255, 255, 0.08);
-}
-
-/* Status container */
-.status-container {
-    margin-bottom: 20px;
-    text-align: center;
-    width: 100%;
-}
-
-/* Response styling */
-.response-container, .response-wrapper {
-    background: rgba(15, 23, 42, 0.6);
-    backdrop-filter: blur(10px);
-    border-radius: 12px;
-    padding: 25px;
-    margin-top: 20px;
-    box-shadow: 0 4px 15px rgba(0, 0, 0, 0.15);
-    max-height: 700px;
-    overflow-y: auto;
-    width: 100%;
-    border: 1px solid rgba(255, 255, 255, 0.05);
-    transition: all 0.3s ease;
-}
-
-.answer-container {
-    background: rgba(30, 41, 59, 0.4);
-    border-radius: 10px;
-    padding: 20px;
-    margin-bottom: 15px;
-    border: 1px solid rgba(255, 255, 255, 0.05);
-}
-
-.answer-container h2 {
-    color: #a5b4fc;
-    margin-top: 25px;
-    margin-bottom: 15px;
-    font-size: 1.5em;
-    border-bottom: 1px solid rgba(255, 255, 255, 0.1);
-    padding-bottom: 10px;
-}
-
-.answer-container h3 {
-    color: #93c5fd;
-    margin-top: 20px;
-    margin-bottom: 12px;
-    font-size: 1.3em;
-}
-
-.answer-container h4 {
-    color: #bae6fd;
-    margin-top: 18px;
-    margin-bottom: 10px;
-    font-size: 1.15em;
-}
-
-.answer-container p {
-    margin-bottom: 15px;
-    line-height: 1.6;
-    color: rgba(255, 255, 255, 0.9);
-    font-size: 1.05em;
-}
-
-/* Status indicators */
-.status-ready {
-    background: linear-gradient(135deg, #134e5e 0%, #71b280 100%);
-    color: white;
-    padding: 12px 18px;
-    border-radius: 30px;
-    margin: 10px auto;
-    font-weight: 600;
-    box-shadow: 0 4px 12px rgba(19, 78, 94, 0.3);
-    display: inline-block;
-    min-width: 180px;
-}
-
-.status-processing {
-    background: linear-gradient(135deg, #2980b9 0%, #2c3e50 100%);
-    color: white;
-    padding: 12px 18px;
-    border-radius: 30px;
-    margin: 10px auto;
-    font-weight: 600;
-    box-shadow: 0 4px 12px rgba(41, 128, 185, 0.3);
-    animation: pulse 2s infinite;
-    display: inline-block;
-    min-width: 180px;
-}
-
-.status-error {
-    background: linear-gradient(135deg, #cb2d3e 0%, #ef473a 100%);
-    color: white;
-    padding: 12px 18px;
-    border-radius: 30px;
-    margin: 15px auto;
-    font-weight: 600;
-    box-shadow: 0 4px 12px rgba(203, 45, 62, 0.3);
-    display: inline-block;
-    min-width: 180px;
-}
-
-.status-complete {
-    background: linear-gradient(135deg, #11998e 0%, #38ef7d 100%);
-    color: white;
-    padding: 12px 18px;
-    border-radius: 30px;
-    margin: 10px auto;
-    font-weight: 600;
-    box-shadow: 0 4px 12px rgba(17, 153, 142, 0.3);
-    display: inline-block;
-    min-width: 180px;
-}
-
-/* Input container */
-.input-container {
-    background: rgba(30, 41, 59, 0.7);
-    backdrop-filter: blur(10px);
-    border-radius: 12px;
-    padding: 25px;
-    margin-bottom: 20px;
-    box-shadow: 0 4px 15px rgba(0, 0, 0, 0.15);
-    border: 1px solid rgba(255, 255, 255, 0.08);
-    transition: all 0.3s ease;
-}
-
-.input-container:focus-within {
-    box-shadow: 0 8px 25px rgba(0, 0, 0, 0.25);
-    transform: translateY(-2px);
-}
-
-.hint-text {
-    color: rgba(255, 255, 255, 0.7);
-    font-size: 0.95em;
-    margin-top: 8px;
-    font-style: italic;
-}
-
-/* Progress container */
-.progress-container {
-    background: rgba(30, 41, 59, 0.7);
-    backdrop-filter: blur(10px);
-    border-radius: 12px;
-    padding: 20px;
-    margin: 20px 0;
-    box-shadow: 0 4px 15px rgba(0, 0, 0, 0.15);
-    width: 100%;
-    min-height: 90px;
-    display: block;
-    border: 1px solid rgba(255, 255, 255, 0.08);
-}
-
-/* Button container */
-.button-container {
-    display: flex;
-    gap: 15px;
-    margin: 25px 0;
-    width: 100%;
-    justify-content: center;
-}
-
-.button-container button {
-    min-height: 52px;
-    font-size: 1.05em;
-    font-weight: 600;
-    padding: 0 35px;
-    flex: 0 1 auto;
-    min-width: 180px;
-    max-width: 220px;
-    letter-spacing: 0.3px;
-    border-radius: 30px !important;
-}
-
-/* Progress label */
-.progress-label {
-    margin-bottom: 12px;
-    color: rgba(255, 255, 255, 0.85);
-    font-weight: 600;
-    font-size: 1.05em;
-    text-align: center;
-}
-
-/* Progress stages styling */
-.progress-stages-container {
-    display: flex;
-    justify-content: space-between;
-    flex-wrap: wrap;
-    gap: 12px;
-    margin-top: 20px;
-    width: 100%;
-}
-
-.progress-stage {
-    padding: 12px 15px;
-    border-radius: 30px;
-    background: rgba(15, 23, 42, 0.7);
-    color: rgba(255, 255, 255, 0.6);
-    font-size: 0.95em;
-    transition: all 0.3s ease;
-    flex: 1;
-    text-align: center;
-    min-width: 120px;
-    max-width: 200px;
-    display: flex;
-    align-items: center;
-    justify-content: center;
-    border: 1px solid rgba(255, 255, 255, 0.05);
-}
-
-.progress-stage.active {
-    background: linear-gradient(135deg, #2980b9 0%, #2c3e50 100%);
-    color: white;
-    font-weight: 600;
-    animation: pulse 2s infinite;
-    border: none;
-    box-shadow: 0 4px 12px rgba(41, 128, 185, 0.25);
-}
-
-.progress-stage.complete {
-    background: linear-gradient(135deg, #11998e 0%, #38ef7d 100%);
-    color: white;
-    font-weight: 600;
-    border: none;
-    box-shadow: 0 4px 12px rgba(17, 153, 142, 0.25);
-}
-
-.progress-stage.error {
-    background: linear-gradient(135deg, #cb2d3e 0%, #ef473a 100%);
-    color: white;
-    font-weight: 600;
-    border: none;
-    box-shadow: 0 4px 12px rgba(203, 45, 62, 0.25);
-}
-
-/* Reasoning steps styling */
-.reasoning-step {
-    background: rgba(15, 23, 42, 0.5);
-    border-radius: 10px;
-    padding: 15px;
-    margin-bottom: 15px;
-    border-left: 3px solid rgba(255, 255, 255, 0.2);
-    box-shadow: 0 2px 10px rgba(0, 0, 0, 0.1);
-}
-
-.reasoning-stage {
-    font-weight: 600;
-    color: #93c5fd;
-    margin-bottom: 10px;
-    padding-bottom: 6px;
-    border-bottom: 1px solid rgba(255, 255, 255, 0.1);
-}
-
-/* Sources section styling */
-.sources-section {
-    background: rgba(15, 23, 42, 0.5);
-    border-radius: 10px;
-    padding: 15px;
-    margin-top: 25px;
-    border-left: 3px solid rgba(255, 255, 255, 0.2);
-    box-shadow: 0 2px 10px rgba(0, 0, 0, 0.1);
-}
 
-.sources-section h3 {
-    color: #93c5fd;
-    margin-top: 0;
-    margin-bottom: 12px;
-}
+# Get the custom CSS
+custom_css = get_custom_css()
 
-.source-item {
-    margin-bottom: 8px;
-    color: rgba(255, 255, 255, 0.85);
-}
-
-/* Enhanced visible styles for inputs */
-input, textarea {
-    background: rgba(15, 23, 42, 0.6) !important;
-    color: white !important;
-    border: 1px solid rgba(255, 255, 255, 0.1) !important;
-    border-radius: 10px !important;
-    padding: 15px !important;
-    font-size: 1.05em !important;
-    transition: all 0.3s ease !important;
-}
-
-input:focus, textarea:focus {
-    border-color: rgba(147, 197, 253, 0.5) !important;
-    outline: none !important;
-    box-shadow: 0 0 0 3px rgba(147, 197, 253, 0.25) !important;
-    background: rgba(15, 23, 42, 0.8) !important;
-}
-
-/* Make buttons more beautiful with gradients */
-button {
-    border-radius: 30px !important;
-    transition: all 0.3s ease !important;
-    outline: none !important;
-    font-weight: 600 !important;
-    letter-spacing: 0.3px !important;
-}
-
-button:hover {
-    transform: translateY(-3px) !important;
-    box-shadow: 0 6px 15px rgba(0, 0, 0, 0.25) !important;
-}
-
-button:active {
-    transform: translateY(-1px) !important;
-}
-
-/* Primary button styling */
-button[variant="primary"], 
-.gr-button-primary {
-    background: linear-gradient(135deg, #4f46e5 0%, #3b82f6 100%) !important;
-    color: white !important;
-    border: none !important;
-}
-
-button[variant="primary"]:hover,
-.gr-button-primary:hover {
-    background: linear-gradient(135deg, #4338ca 0%, #2563eb 100%) !important;
-}
-
-/* Secondary button styling */
-button[variant="secondary"],
-.gr-button-secondary {
-    background: rgba(30, 41, 59, 0.8) !important;
-    color: white !important;
-    border: 1px solid rgba(255, 255, 255, 0.15) !important;
-}
-
-button[variant="secondary"]:hover,
-.gr-button-secondary:hover {
-    background: rgba(30, 41, 59, 0.95) !important;
-}
-
-/* Animation for processing effects */
-@keyframes pulse {
-    0% {
-        opacity: 1;
-        box-shadow: 0 4px 12px rgba(41, 128, 185, 0.3);
-    }
-    50% {
-        opacity: 0.85;
-        box-shadow: 0 4px 20px rgba(41, 128, 185, 0.5);
-    }
-    100% {
-        opacity: 1;
-        box-shadow: 0 4px 12px rgba(41, 128, 185, 0.3);
-    }
-}
-
-/* Responsive improvements */
-@media (min-width: 768px) {
-    .gradio-container {
-        min-width: 750px !important;
-        max-width: 1200px !important;
-    }
-    
-    .input-container, .output-container {
-        padding: 30px !important;
-    }
-    
-    .button-container {
-        justify-content: center !important;
-        gap: 25px !important;
-    }
-}
-
-@media (max-width: 767px) {
-    .progress-stage {
-        min-width: 90px;
-        padding: 10px;
-        font-size: 0.9em;
-    }
-    
-    .app-header h1 {
-        font-size: 1.8rem;
-    }
-    
-    .app-header p {
-        font-size: 1rem;
-    }
-    
-    .logo-icon {
-        font-size: 1.5rem;
-    }
-    
-    .logo-text {
-        font-size: 1.2rem;
-    }
-}
-
-/* Override any built-in Gradio container colors */
-.gradio-container {
-    color: white !important;
-}
-
-.gr-box, .gr-form, .gr-panel {
-    background: rgba(30, 41, 59, 0.7) !important;
-    backdrop-filter: blur(10px) !important;
-    color: white !important;
-    border-radius: 12px !important;
-    border: 1px solid rgba(255, 255, 255, 0.08) !important;
-}
-
-.gr-input, .gr-textarea, .gr-checkbox, .gr-radio, .gr-slider, .gr-dropdown {
-    background: rgba(15, 23, 42, 0.6) !important;
-    color: white !important;
-    border-color: rgba(255, 255, 255, 0.1) !important;
-}
-
-/* Override labels and markdown content */
-.gr-label, .gr-markdown {
-    color: rgba(255, 255, 255, 0.9) !important;
-}
-
-/* Make sure accordions are styled properly */
-.gr-accordion {
-    background: rgba(30, 41, 59, 0.7) !important;
-    backdrop-filter: blur(10px) !important;
-    color: white !important;
-    border: 1px solid rgba(255, 255, 255, 0.08) !important;
-    border-radius: 12px !important;
-    margin: 15px 0 !important;
-}
-
-.gr-accordion-header {
-    background: rgba(15, 23, 42, 0.5) !important;
-    color: white !important;
-    font-weight: 600 !important;
-    padding: 15px !important;
-    border-radius: 12px 12px 0 0 !important;
-}
-
-.gr-accordion-content {
-    background: transparent !important;
-    color: white !important;
-    padding: 20px !important;
-}
-
-/* Ensure response text is visible */
-.response-wrapper, 
-.response-wrapper *,
-.answer-container, 
-.answer-container * {
-    color: rgba(255, 255, 255, 0.9) !important;
-    background-color: transparent !important;
-}
-
-.response-wrapper {
-    background: rgba(15, 23, 42, 0.6) !important;
-    backdrop-filter: blur(10px) !important;
-    border-radius: 12px !important;
-    padding: 25px !important;
-    margin-top: 20px !important;
-    box-shadow: 0 4px 15px rgba(0, 0, 0, 0.15) !important;
-    border: 1px solid rgba(255, 255, 255, 0.05) !important;
-}
-
-.answer-container {
-    background: rgba(30, 41, 59, 0.4) !important;
-    border-radius: 10px !important;
-    padding: 20px !important;
-    margin-bottom: 15px !important;
-    border: 1px solid rgba(255, 255, 255, 0.05) !important;
-}
-
-/* Apply styles to specific elements within the answer container */
-.answer-container h2 {
-    color: #a5b4fc !important;
-    margin-top: 25px !important;
-    margin-bottom: 15px !important;
-    font-size: 1.5em !important;
-    border-bottom: 1px solid rgba(255, 255, 255, 0.1) !important;
-    padding-bottom: 10px !important;
-}
-
-.answer-container h3 {
-    color: #93c5fd !important;
-    margin-top: 20px !important;
-    margin-bottom: 12px !important;
-    font-size: 1.3em !important;
-}
-
-.answer-container h4 {
-    color: #bae6fd !important;
-    margin-top: 18px !important;
-    margin-bottom: 10px !important;
-    font-size: 1.15em !important;
-}
-
-.answer-container p {
-    margin-bottom: 15px !important;
-    line-height: 1.6 !important;
-    color: rgba(255, 255, 255, 0.9) !important;
-    font-size: 1.05em !important;
-}
-
-.answer-container ul,
-.answer-container ol {
-    margin-left: 20px !important;
-    margin-bottom: 15px !important;
-    color: rgba(255, 255, 255, 0.9) !important;
-}
-
-.answer-container li {
-    margin-bottom: 8px !important;
-    color: rgba(255, 255, 255, 0.9) !important;
-    line-height: 1.5 !important;
-}
-
-/* Fix raw HTML rendering */
-.prose {
-    color: rgba(255, 255, 255, 0.9) !important;
-    background-color: transparent !important;
-}
-
-.prose * {
-    color: inherit !important;
-}
-
-/* Slider styling */
-.gr-slider {
-    background: rgba(15, 23, 42, 0.6) !important;
-}
-
-.gr-slider .thumb {
-    background: #4f46e5 !important;
-    border: 2px solid white !important;
-}
-
-.gr-slider .track {
-    background: rgba(255, 255, 255, 0.2) !important;
-}
-
-.gr-slider .track-fill {
-    background: linear-gradient(90deg, #4f46e5 0%, #3b82f6 100%) !important;
-}
-
-/* Checkbox styling */
-.gr-checkbox {
-    border: 2px solid rgba(255, 255, 255, 0.3) !important;
-    background: rgba(15, 23, 42, 0.6) !important;
-}
-
-.gr-checkbox:checked {
-    background: #4f46e5 !important;
-    border-color: white !important;
-}
-
-/* Better scrollbars for the response area */
-.response-wrapper::-webkit-scrollbar,
-.answer-container::-webkit-scrollbar {
-    width: 8px;
-    height: 8px;
-}
-
-.response-wrapper::-webkit-scrollbar-track,
-.answer-container::-webkit-scrollbar-track {
-    background: rgba(255, 255, 255, 0.05);
-    border-radius: 4px;
-}
-
-.response-wrapper::-webkit-scrollbar-thumb,
-.answer-container::-webkit-scrollbar-thumb {
-    background: rgba(255, 255, 255, 0.2);
-    border-radius: 4px;
-}
-
-.response-wrapper::-webkit-scrollbar-thumb:hover,
-.answer-container::-webkit-scrollbar-thumb:hover {
-    background: rgba(255, 255, 255, 0.3);
-}
-"""
-
-# Create the Gradio interface with full functionality
-def create_interface():
-    """Create the Gradio interface."""
-    
-    with gr.Blocks(
-        title="Defensor Legal Assistant",
-        theme=gr.themes.Base(),
-        css=custom_css
-    ) as demo:
-        # Create a header with app-header class for better styling
-        with gr.Column(elem_classes="app-header"):
-            gr.HTML("""
-            <div class="logo-container">
-                <div class="logo">
-                    <span class="logo-icon">⚖️</span>
-                    <span class="logo-text">Defensor</span>
-                </div>
-            </div>
-            """)
-            gr.Markdown("# Defensor Legal Assistant")
-            gr.Markdown("Ask questions about legal defense strategies and get comprehensive answers based on legitimate sources.")
-        
-        # Main application container
-        with gr.Column(elem_classes="main-container"):
-            # Add status indicator with styles - in its own container
-            with gr.Row(elem_classes="status-container"):
-                status_indicator = gr.HTML(
-                    value="<div class='status-ready'>Ready to assist</div>",
-                    label=None
-                )
-            
-            # Enhanced container for the input area
-            with gr.Column(elem_classes="input-container"):
-                query_input = gr.Textbox(
-                    label="Your Legal Question",
-                    placeholder="Enter your legal question here (e.g., 'What are the key elements of self-defense?')",
-                    lines=3
-                )
-                gr.Markdown("<div class='hint-text'>Ask detailed questions for better results. Be specific about the legal concepts you're interested in.</div>")
-            
-            # Create a dedicated container for the progress area
-            with gr.Column(elem_classes="progress-container"):
-                # Progress indicator with animation and better styling
-                progress = gr.Progress()
-                
-                # Add a visual progress label
-                progress_label = gr.HTML(
-                    "<div class='progress-label'>Waiting for your question...</div>", 
-                    visible=True
-                )
-            
-            # Better button container with centered buttons
-            with gr.Row(elem_classes="button-container"):
-                submit_btn = gr.Button(
-                    "Submit Question", 
-                    variant="primary", 
-                    interactive=True
-                )
-                clear_btn = gr.Button(
-                    "Clear", 
-                    variant="secondary"
-                )
-            
-            # Settings in collapsible section
-            with gr.Accordion("Settings & Options", open=False):
-                with gr.Row():
-                    top_k = gr.Slider(
-                        minimum=5, 
-                        maximum=50, 
-                        value=20, 
-                        step=5, 
-                        label="Number of Documents to Consider"
-                    )
-                    debug = gr.Checkbox(label="Show Agent Reasoning Process", value=True)
-            
-            # Display area for the response with enhanced formatting
-            with gr.Column(elem_classes="output-container"):
-                output = gr.HTML(label="Response", value="")
-        
-        # About section at the bottom
-        with gr.Accordion("About Defensor Legal Assistant", open=False):
-            gr.Markdown("""
-            ## About Defensor Legal Assistant
-            
-            Defensor is a specialized legal assistant powered by advanced AI technology and a comprehensive legal knowledge base. It provides insights and information about legal defense strategies based on authoritative sources.
-            
-            ### How It Works
-            
-            1. **Query Analysis**: The system analyzes your legal question to understand its context and intent
-            2. **Knowledge Retrieval**: It searches through a database of legal documents to find relevant information
-            3. **Context Organization**: The system organizes the retrieved information into a coherent context
-            4. **Response Generation**: A comprehensive answer is generated based on the collected information
-            
-            ### Tips for Best Results
-            
-            - Be specific in your questions
-            - Include relevant legal terms or concepts
-            - For complex topics, consider breaking down into multiple targeted questions
-            - Use the "Show Agent Reasoning" option to see how the system arrived at its answer
-            
-            ### Limitations
-            
-            This assistant provides legal information, not legal advice. Always consult with a qualified attorney for advice on specific legal matters.
-            """)
-        
-        def process_query_with_status(query, top_k=50, debug=False, progress=gr.Progress()):
-            """Process a query using the agent director with status updates."""
-            global director
-            
-            if not query.strip():
-                return "<div class='status-error'>Please enter a question to proceed</div>", ""
-                
-            # Show progress - with fixed percentage formatting
-            progress(0, desc="Initializing...")
-            
-            # Import here to avoid circular references
-            try:
-                from src.agents.agent_director import AgentDirector
-                from src.utils.config import CHAT_MODEL
-            except ImportError as e:
-                return f"<div class='status-error'>ERROR: Failed to import required modules: {str(e)}</div>", ""
-            
-            # Initialize the agent if not already done
-            if director is None:
-                try:
-                    progress(0.1, desc="Setting up the agent...")  # Use decimal for percentage
-                    # Initialize the agent director
-                    director = AgentDirector(model=CHAT_MODEL, top_k=top_k, debug=debug)
-                    progress(0.2, desc="Agent ready")  # Use decimal for percentage
-                except Exception as e:
-                    error_details = traceback.format_exc()
-                    print(f"Agent initialization error: {e}")
-                    print(error_details)
-                    progress(1.0, desc="Error")  # Use decimal for percentage
-                    return f"<div class='status-error'>ERROR: Failed to initialize agent: {str(e)}</div>", ""
-            
-            # Process the query
-            try:
-                # Update progress with meaningful steps and show visual stages
-                progress_stages = """<div class='progress-stages-container'>
-                    <div class='progress-stage active'>► Analyzing query</div>
-                    <div class='progress-stage'>Searching documents</div>
-                    <div class='progress-stage'>Organizing information</div>
-                    <div class='progress-stage'>Formulating response</div>
-                </div>"""
-                
-                # Update progress with meaningful steps
-                progress(0.25, desc="Analyzing your question...")  # Use decimal for percentage
-                time.sleep(0.5)  # Small delay for visual feedback
-                
-                progress(0.4, desc="Searching knowledge base...")  # Use decimal for percentage
-                progress_stages = """<div class='progress-stages-container'>
-                    <div class='progress-stage complete'>✓ Analyzing query</div>
-                    <div class='progress-stage active'>► Searching documents</div>
-                    <div class='progress-stage'>Organizing information</div>
-                    <div class='progress-stage'>Formulating response</div>
-                </div>"""
-                time.sleep(0.5)  # Small delay for visual feedback
-                
-                progress(0.6, desc="Organizing information...")  # Use decimal for percentage
-                progress_stages = """<div class='progress-stages-container'>
-                    <div class='progress-stage complete'>✓ Analyzing query</div>
-                    <div class='progress-stage complete'>✓ Searching documents</div>
-                    <div class='progress-stage active'>► Organizing information</div>
-                    <div class='progress-stage'>Formulating response</div>
-                </div>"""
-                time.sleep(0.5)  # Small delay for visual feedback
-                
-                progress(0.8, desc="Formulating legal response...")  # Use decimal for percentage
-                progress_stages = """<div class='progress-stages-container'>
-                    <div class='progress-stage complete'>✓ Analyzing query</div>
-                    <div class='progress-stage complete'>✓ Searching documents</div>
-                    <div class='progress-stage complete'>✓ Organizing information</div>
-                    <div class='progress-stage active'>► Formulating response</div>
-                </div>"""
-                
-                # Process the query with timeout handling
-                try:
-                    # Set a timeout for processing
-                    import threading
-                    import queue
-
-                    def process_with_timeout():
-                        try:
-                            result = director.process_query(query)
-                            result_queue.put(result)
-                        except Exception as e:
-                            result_queue.put({"error": str(e)})
-
-                    result_queue = queue.Queue()
-                    processing_thread = threading.Thread(target=process_with_timeout)
-                    processing_thread.daemon = True
-                    processing_thread.start()
-                    
-                    # Wait for the result with a timeout
-                    timeout = 60  # 60 seconds timeout
-                    processing_thread.join(timeout)
-                    
-                    if not result_queue.empty():
-                        result = result_queue.get()
-                        if "error" in result:
-                            raise Exception(result["error"])
-                    else:
-                        # If the queue is empty after timeout, provide a fallback answer
-                        print("Query processing timed out, providing fallback response")
-                        
-                        # Try to use the legal agent as fallback
-                        try:
-                            from src.agents.legal_agent import LegalAgent
-                            fallback_agent = LegalAgent(model=CHAT_MODEL)
-                            # Let the legal agent handle retrieval internally
-                            result = fallback_agent.answer_query(query, top_k=5)
-                        except Exception as fallback_error:
-                            print(f"Fallback error: {fallback_error}")
-                            # If fallback fails, use a generic response
-                            result = {
-                                "answer": "I apologize, but I was unable to generate a complete response. Please try rephrasing your question or asking something else.",
-                                "sources": []
-                            }
-                except Exception as process_error:
-                    print(f"Error during query processing: {process_error}")
-                    traceback.print_exc()
-                    
-                    # Provide a simplified fallback
-                    result = {
-                        "answer": f"I encountered an error while processing your query: {str(process_error)}. Please try again or rephrase your question.",
-                        "sources": []
-                    }
-                
-                # Format the response with enhanced HTML
-                answer = result.get("answer", "No answer available.")
-                
-                # Create a well-formatted, dark-theme friendly output
-                answer_html = ""
-                
-                # Convert answer text to HTML paragraphs
-                for paragraph in answer.split("\n\n"):
-                    if paragraph.strip():
-                        # Check if it's a heading
-                        if paragraph.startswith("# "):
-                            heading = paragraph[2:].strip()
-                            answer_html += f"<h2>{heading}</h2>"
-                        elif paragraph.startswith("## "):
-                            heading = paragraph[3:].strip()
-                            answer_html += f"<h3>{heading}</h3>"
-                        elif paragraph.startswith("### "):
-                            heading = paragraph[4:].strip()
-                            answer_html += f"<h4>{heading}</h4>"
-                        else:
-                            answer_html += f"<p>{paragraph}</p>"
-                
-                # Add sources if available
-                if "sources" in result and result["sources"]:
-                    sources_html = "<div class='sources-section'>"
-                    sources_html += "<h3>Sources</h3><ul>"
-                    for source in result["sources"]:
-                        if source and source != "Source information unavailable":
-                            sources_html += f"<li class='source-item'>{source}</li>"
-                    sources_html += "</ul></div>"
-                    answer_html += sources_html
-                
-                # Add debugging information if available with nice formatting
-                if debug and result.get("reasoning_steps") is not None:
-                    reasoning_html = "<h3>Agent Reasoning Process</h3>"
-                    for step in result["reasoning_steps"]:
-                        step_stage = step['stage']
-                        step_reasoning = step['reasoning'].replace('\n', '<br>')
-                        reasoning_html += f"""
-                        <div class='reasoning-step'>
-                            <div class='reasoning-stage'>{step_stage}</div>
-                            <div>{step_reasoning}</div>
-                        </div>
-                        """
-                    answer_html += reasoning_html
-                
-                progress(1.0, desc="Completed!")
-                
-                # Update the progress stages to show completion
-                progress_stages = """<div class='progress-stages-container'>
-                    <div class='progress-stage complete'>✓ Analyzing query</div>
-                    <div class='progress-stage complete'>✓ Searching documents</div>
-                    <div class='progress-stage complete'>✓ Organizing information</div>
-                    <div class='progress-stage complete'>✓ Response ready</div>
-                </div>"""
-                
-                # Final formatted output with all styling embedded
-                return f"""
-                <div class="response-wrapper">
-                    <div class="answer-container">
-                        {answer_html}
-                    </div>
-                </div>
-                """, progress_stages
-            
-            except Exception as e:
-                error_details = traceback.format_exc()
-                print(f"Query processing error: {e}")
-                print(error_details)
-                progress(1.0, desc="Error")
-                
-                # Create a complete progress stages display with error indicator
-                error_progress_stages = """<div class='progress-stages-container'>
-                    <div class='progress-stage complete'>✓ Analyzing query</div>
-                    <div class='progress-stage complete'>✓ Searching documents</div>
-                    <div class='progress-stage error'>✗ Error</div>
-                    <div class='progress-stage'>Formulating response</div>
-                </div>"""
-                
-                return f"""
-                <div class='status-error'>
-                    <strong>ERROR:</strong> {str(e)}
-                    <p>Please try again or rephrase your question.</p>
-                </div>
-                """, error_progress_stages
-        
-        def set_ready_status():
-            return "<div class='status-ready'>Ready to assist</div>", "<div class='progress-label'>Waiting for your question...</div>"
-        
-        def set_processing_status():
-            return "<div class='status-processing'>Processing your question...</div>", gr.update(interactive=False), "<div class='progress-label'>Processing legal query...</div>"
-        
-        def set_completed_status():
-            return "<div class='status-complete'>Response complete</div>", gr.update(interactive=True), "<div class='progress-label'>Response complete</div>"
-        
-        def clear_inputs():
-            return "", "<div class='status-ready'>Ready to assist</div>", gr.update(interactive=True), "", "<div class='progress-label'>Waiting for your question...</div>"
-        
-        # Set up event handlers with improved button state management
-        submit_btn.click(
-            fn=set_processing_status,
-            outputs=[status_indicator, submit_btn, progress_label],
-            queue=False
-        ).then(
-            fn=process_query_with_status,
-            inputs=[query_input, top_k, debug],
-            outputs=[output, progress_label],
-            api_name="query"
-        ).then(
-            fn=set_completed_status,
-            outputs=[status_indicator, submit_btn, progress_label],
-            queue=False
-        )
-        
-        clear_btn.click(
-            fn=clear_inputs,
-            outputs=[query_input, status_indicator, submit_btn, output, progress_label],
-            queue=False
-        )
-        
-        return demo
+# Define a wrapper function for processing queries
+def process_query_wrapper(query, top_k=50, debug=False, progress=None):
+    """Wrapper function to process queries with the resource manager."""
+    try:
+        # Get the director from the resource manager
+        director = resource_manager.get_director()
+        return process_query_with_status(query, top_k, debug, progress, director)
+    except Exception as e:
+        import traceback
+        print(f"Error in process_query_wrapper: {e}")
+        traceback.print_exc()
+        return (f"""
+        <div class='status-error'>
+            <strong>ERROR:</strong> An unexpected error occurred: {str(e)}
+            <p>Please try again or contact support if the problem persists.</p>
+        </div>
+        """, "")
 
 # Launch the app
 if __name__ == "__main__":
     # Create and launch the appropriate interface based on data readiness
     if data_ready:
-        demo = create_interface()
+        demo = create_interface(custom_css, process_query_wrapper)
     else:
-        demo = create_fallback_ui()
+        demo = create_fallback_ui(custom_css)
     
     demo.launch(
         server_name="0.0.0.0",