import os import sys import signal import warnings import gradio as gr import config import ingestion import vector_store import rag_engine import logic # Suppress warnings warnings.filterwarnings("ignore", category=FutureWarning) # Global variable to store the chain rag_chain = None def initialize_system_once(): """ Initialize the complete system only once. """ global rag_chain if rag_chain is not None: return rag_chain print("Initializing LegalizeAI System...") # Initialize Embedding Model print("Loading embedding model...") embedding_model = vector_store.get_embedding_model() # Check if Vector Store exists if os.path.exists(config.CHROMA_DB_DIR) and os.listdir(config.CHROMA_DB_DIR): print(f"Loading existing vector store from {config.CHROMA_DB_DIR}...") v_store = vector_store.get_vector_store(embedding_model) else: print("No existing vector store found. Starting ingestion process...") docs = ingestion.load_documents() if not docs: # Create empty placeholder if no docs found to prevent crash, # but warn hard. print("CRITICAL WARNING: No documents loaded. App will run but local search will fail.") # In a real app we might want to fail, but for UI it's better to stay up return None chunks = ingestion.split_documents(docs) print("Creating vector store...") v_store = vector_store.create_vector_store(chunks, embedding_model) # Setup Retriever retriever = vector_store.get_retriever(v_store) # Setup RAG Chain print("Initializing RAG chain...") rag_chain = rag_engine.create_rag_chain(retriever) print("System initialization complete!") return rag_chain def chat_response(message, history): """ Gradio chat function. """ try: chain = initialize_system_once() if not chain: return "System Error: Failed to initialize AI chain. Please check server logs." response = logic.generate_hybrid_response(message, chain) return response except Exception as e: return f"An error occurred: {str(e)}" # Custom CSS for a professional look custom_css = """ body { background-color: #f0f2f5; } footer { visibility: hidden !important; } /* Custom Developer Footer */ .dev-footer { text-align: center; padding: 20px; margin-top: 30px; border-top: 1px solid #e5e7eb; color: #4b5563; background-color: transparent !important; } .dev-footer a { display: inline-flex; align-items: center; justify-content: center; margin: 0 10px; color: #4b5563; text-decoration: none; transition: color 0.2s; } .dev-footer a:hover { color: #1f2937; } .dev-footer svg { margin-right: 5px; width: 20px; height: 20px; fill: currentColor; } /* Chatbot styling */ #chatbot { min-height: 500px; max-height: 500px; overflow-y: auto; } /* Mobile responsiveness */ @media (max-width: 768px) { .dev-footer div { flex-direction: column; gap: 10px; } .dev-footer a { margin: 5px 0; } #chatbot { min-height: 400px; max-height: 400px; } } """ # HTML for the footer footer_html = """ """ # Create the Gradio Interface def create_ui(): initialize_system_once() # Define interaction functions def interact(message, history): if not message: return "", history # Initialize history if strictly None (though usually empty list) if history is None: history = [] # Add user message history.append({"role": "user", "content": message}) # Get response try: response = chat_response(message, history) except Exception as e: response = f"Error: {str(e)}" # Add assistant response history.append({"role": "assistant", "content": response}) return "", history def retry_last(history): if not history: return history, "" # Pop last message if it's assistant if history and history[-1]["role"] == "assistant": history.pop() # Pop user message to edit if history and history[-1]["role"] == "user": last_msg = history.pop() return history, last_msg["content"] return history, "" def regenerate_response(history): """Regenerate the last assistant response.""" if not history or len(history) < 2: return history # Get the last user message last_user_msg = None for msg in reversed(history): if msg["role"] == "user": last_user_msg = msg["content"] break if last_user_msg: # Remove the last assistant response if history[-1]["role"] == "assistant": history.pop() # Generate new response try: response = chat_response(last_user_msg, history) history.append({"role": "assistant", "content": response}) except Exception as e: history.append({"role": "assistant", "content": f"Error: {str(e)}"}) return history # Create the Gradio Blocks with gr.Blocks(title="⚖️ LegalizeAI") as demo: gr.Markdown( """ # ⚖️ LegalizeAI **Professional Assistant for Pakistani Law** Consulting Constitution of Pakistan, Pakistan Penal Code, and Real-time Web Sources. *Disclaimer: This AI provides legal information but is not a substitute for professional legal advice.* """ ) chatbot = gr.Chatbot( height=500, elem_id="chatbot", avatar_images=(None, "⚖️") # Removed show_copy_button for compatibility with older Gradio versions ) with gr.Row(): txt = gr.Textbox( scale=4, show_label=False, placeholder="Ask a legal question...", container=False, autofocus=True, max_lines=3 ) submit_btn = gr.Button("Send 🚀", scale=1, variant="primary") with gr.Row(): retry_btn = gr.Button("Retry 🔄", size="sm") clear_btn = gr.Button("Clear 🗑️", size="sm") regenerate_btn = gr.Button("Regenerate 🔁", size="sm") # Example buttons logic examples = [ "What is the punishment for theft in Pakistan?", "Explain Article 62 of the Constitution.", "Who is the current Prime Minister?", "What are my fundamental rights?", "Explain the process of filing a criminal case.", "What is defamation under Pakistani law?" ] gr.Examples( examples=examples, inputs=txt, examples_per_page=3 ) # Footer gr.HTML(footer_html) # Event Wiring submit_btn.click(interact, [txt, chatbot], [txt, chatbot]) txt.submit(interact, [txt, chatbot], [txt, chatbot]) retry_btn.click(retry_last, [chatbot], [chatbot, txt]) clear_btn.click(lambda: None, None, chatbot, queue=False) regenerate_btn.click(regenerate_response, [chatbot], [chatbot]) return demo def cleanup(signum, frame): """Cleanup function for graceful shutdown.""" print("\n🛑 Shutting down LegalizeAI gracefully...") sys.exit(0) def main(): """Main entry point.""" try: # Register cleanup handlers signal.signal(signal.SIGINT, cleanup) signal.signal(signal.SIGTERM, cleanup) # Initialize system once before UI creation print("🚀 Starting LegalizeAI...") initialize_system_once() # Determine if running on Hugging Face Spaces is_huggingface = os.getenv("SPACE_ID") is not None server_name = "0.0.0.0" if is_huggingface else "127.0.0.1" if is_huggingface: print("🌐 Running on Hugging Face Spaces") else: print("💻 Running locally") # Using a professional theme theme = gr.themes.Soft( primary_hue="slate", secondary_hue="stone", neutral_hue="zinc", font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"] ) # Create UI demo = create_ui() # Launch with appropriate settings print(f"🌍 Server starting on {server_name}:7860") demo.launch( server_name=server_name, server_port=7860, share=False, # Don't use share=True on Spaces favicon_path=None, theme=theme, show_error=True, quiet=False, show_api=False, css=custom_css # CSS moved here for Gradio 6.0+ compatibility ) except KeyboardInterrupt: print("\n👋 Shutdown requested by user") sys.exit(0) except Exception as e: print(f"❌ Fatal Error: {e}") import traceback traceback.print_exc() sys.exit(1) if __name__ == "__main__": # Set environment variables for better performance os.environ["TOKENIZERS_PARALLELISM"] = "false" os.environ["PYTHONWARNINGS"] = "ignore" # Run the application main()