Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from config import ( | |
| APP_SUBTITLE, | |
| APP_TITLE, | |
| DEFAULT_CONCURRENCY_LIMIT, | |
| DEFAULT_JURISDICTION, | |
| JURISDICTIONS, | |
| MAX_RETRIEVAL_LOOPS, | |
| TOP_K_RETRIEVAL, | |
| UPLOAD_SESSION_ROOT, | |
| ) | |
| from document_service import ( | |
| clear_document_session, | |
| process_uploaded_document, | |
| validate_document_session, | |
| ) | |
| from pipeline import run_document_rag, run_rag_agent | |
| from job_service import ( | |
| forget_rag_job, | |
| get_rag_job_status, | |
| start_document_processing_job, | |
| start_document_rag_job, | |
| start_rag_job, | |
| ) | |
| from services import check_system_status | |
| CUSTOM_CSS = """ | |
| .gradio-container {max-width:1280px!important;margin:0 auto!important; | |
| font-family:"Segoe UI",Tahoma,Geneva,Verdana,sans-serif!important} | |
| .app-header {text-align:center;padding:24px; | |
| background:linear-gradient(135deg,#5f9eea 0%,#7bc5ef 100%); | |
| color:white;border-radius:10px;margin-bottom:18px} | |
| .app-header h1 {margin:0;font-size:2.35rem;font-weight:650} | |
| .app-header p {margin:8px 0 0;font-size:1.05rem;opacity:.95} | |
| .app-section {background:white;padding:18px;border-radius:10px; | |
| box-shadow:0 3px 12px rgba(0,0,0,.07);margin-bottom:16px} | |
| .answer-panel {border-left:4px solid #5f9eea} | |
| .process-panel {border-left:4px solid #7a8ca5} | |
| .source-panel {border-left:4px solid #65a97b} | |
| .footer {text-align:center;padding:14px;margin-top:24px;color:#666; | |
| font-size:.9rem;border-top:1px solid #e6e6e6} | |
| """ | |
| JURISDICTION_CHOICES = [ | |
| (str(details["name"]), code) | |
| for code, details in JURISDICTIONS.items() | |
| ] | |
| def clear_interface(): | |
| return "", "", "", "", "", "Status: Ready" | |
| with gr.Blocks(title=APP_TITLE, fill_width=True) as demo: | |
| with gr.Column(elem_classes="app-header"): | |
| gr.Markdown(f"# {APP_TITLE}") | |
| gr.Markdown(APP_SUBTITLE) | |
| with gr.Column(elem_classes="app-section"): | |
| gr.Markdown( | |
| "Select a jurisdiction, then enter a nutrient-policy, " | |
| "wastewater-management, or regulatory-compliance question. " | |
| "The agent will search the corresponding regulatory vector " | |
| "database and generate a grounded answer." | |
| ) | |
| with gr.Row(): | |
| jurisdiction_dropdown = gr.Dropdown( | |
| choices=JURISDICTION_CHOICES, | |
| value=DEFAULT_JURISDICTION, | |
| label="Jurisdiction", | |
| interactive=True, | |
| ) | |
| question_box = gr.Textbox( | |
| label="Question", | |
| placeholder=( | |
| "Example: What phosphorus-related compliance requirements " | |
| "apply to municipal wastewater treatment facilities?" | |
| ), | |
| lines=4, | |
| ) | |
| with gr.Row(): | |
| top_k_slider = gr.Slider( | |
| 1, | |
| 15, | |
| value=TOP_K_RETRIEVAL, | |
| step=1, | |
| label="Passages retrieved per search", | |
| ) | |
| max_loops_slider = gr.Slider( | |
| 1, | |
| 10, | |
| value=min(MAX_RETRIEVAL_LOOPS, 10), | |
| step=1, | |
| label="Maximum retrieval steps", | |
| ) | |
| with gr.Row(): | |
| run_button = gr.Button( | |
| "Run RAG Agent", | |
| variant="primary", | |
| ) | |
| clear_button = gr.Button("Clear") | |
| status_button = gr.Button("Check System") | |
| system_status = gr.Textbox( | |
| value="Status: Ready", | |
| label="System status", | |
| lines=20, | |
| interactive=False, | |
| ) | |
| with gr.Column(elem_classes="app-section"): | |
| gr.Markdown("## Uploaded document mode") | |
| gr.Markdown( | |
| "Process one document independently from the Florida and New Jersey " | |
| "regulatory databases. Supported formats are PDF, DOC, DOCX, PPTX, " | |
| "HTML, and TXT. The temporary session expires after 30 minutes " | |
| "of inactivity." | |
| ) | |
| uploaded_document = gr.File( | |
| label="Document", | |
| file_types=[".pdf", ".doc", ".docx", ".pptx", ".html", ".txt"], | |
| type="filepath", | |
| ) | |
| process_document_button = gr.Button( | |
| "Process document", | |
| variant="secondary", | |
| ) | |
| document_requested_session_id = gr.Textbox( | |
| label="Requested document session ID", | |
| visible=False, | |
| ) | |
| document_session_id = gr.Textbox( | |
| label="Document session ID", | |
| visible=False, | |
| ) | |
| document_ready = gr.Checkbox( | |
| label="Document database ready", | |
| value=False, | |
| visible=False, | |
| ) | |
| document_session_valid = gr.Checkbox( | |
| label="Document session valid", | |
| value=False, | |
| visible=False, | |
| ) | |
| document_session_validation_message = gr.Textbox( | |
| label="Document session validation message", | |
| visible=False, | |
| ) | |
| validate_document_session_button = gr.Button( | |
| "Validate Uploaded Document Session", | |
| visible=False, | |
| ) | |
| document_status = gr.Markdown( | |
| "No document has been processed." | |
| ) | |
| vector_download = gr.File( | |
| label="Download temporary vector database", | |
| interactive=False, | |
| ) | |
| document_question = gr.Textbox( | |
| label="Question about the uploaded document", | |
| lines=3, | |
| ) | |
| with gr.Row(): | |
| run_document_button = gr.Button( | |
| "Ask Uploaded Document", | |
| variant="primary", | |
| ) | |
| clear_document_button = gr.Button( | |
| "Clear Uploaded Document Session" | |
| ) | |
| # Hidden components used by the resumable job API. | |
| resumable_requested_job_id = gr.Textbox( | |
| visible=False, | |
| label="Requested resumable job ID", | |
| ) | |
| resumable_job_id = gr.Textbox( | |
| visible=False, | |
| label="Resumable job ID", | |
| ) | |
| resumable_job_status = gr.Textbox( | |
| visible=False, | |
| label="Resumable job status JSON", | |
| ) | |
| resumable_start_button = gr.Button( | |
| "Start resumable state RAG job", | |
| visible=False, | |
| ) | |
| resumable_start_document_button = gr.Button( | |
| "Start resumable document RAG job", | |
| visible=False, | |
| ) | |
| resumable_start_document_processing_button = gr.Button( | |
| "Start resumable document-processing job", | |
| visible=False, | |
| ) | |
| resumable_status_button = gr.Button( | |
| "Get resumable RAG job status", | |
| visible=False, | |
| ) | |
| resumable_forget_button = gr.Button( | |
| "Forget resumable RAG job", | |
| visible=False, | |
| ) | |
| with gr.Column( | |
| elem_classes=["app-section", "answer-panel"] | |
| ): | |
| gr.Markdown("## Grounded answer") | |
| answer_output = gr.Markdown( | |
| "The final grounded answer will appear here." | |
| ) | |
| with gr.Accordion( | |
| "Retrieval process", | |
| open=True, | |
| elem_classes=["app-section", "process-panel"], | |
| ): | |
| process_output = gr.Markdown( | |
| "No retrieval has been run." | |
| ) | |
| with gr.Accordion( | |
| "Retrieved source passages", | |
| open=False, | |
| elem_classes=["app-section", "source-panel"], | |
| ): | |
| sources_output = gr.Markdown( | |
| "No sources have been retrieved." | |
| ) | |
| with gr.Accordion("Run diagnostics", open=False): | |
| diagnostics_output = gr.Textbox( | |
| label="Diagnostics", | |
| lines=12, | |
| interactive=False, | |
| ) | |
| with gr.Column(elem_classes="footer"): | |
| gr.Markdown( | |
| "WNP-ARAG | Built with Gradio, FAISS, " | |
| "LangChain, and Llama 3.1" | |
| ) | |
| inputs = [ | |
| question_box, | |
| jurisdiction_dropdown, | |
| top_k_slider, | |
| max_loops_slider, | |
| ] | |
| outputs = [ | |
| answer_output, | |
| process_output, | |
| sources_output, | |
| diagnostics_output, | |
| ] | |
| run_button.click( | |
| run_rag_agent, | |
| inputs, | |
| outputs, | |
| api_name="run_rag_agent", | |
| api_visibility="public", | |
| ) | |
| resumable_start_button.click( | |
| start_rag_job, | |
| inputs=[ | |
| question_box, | |
| jurisdiction_dropdown, | |
| top_k_slider, | |
| max_loops_slider, | |
| resumable_requested_job_id, | |
| ], | |
| outputs=resumable_job_id, | |
| api_name="start_rag_job", | |
| api_visibility="public", | |
| queue=False, | |
| ) | |
| resumable_start_document_processing_button.click( | |
| start_document_processing_job, | |
| inputs=[ | |
| uploaded_document, | |
| document_requested_session_id, | |
| resumable_requested_job_id, | |
| ], | |
| outputs=resumable_job_id, | |
| api_name="start_document_processing_job", | |
| api_visibility="public", | |
| queue=False, | |
| ) | |
| resumable_start_document_button.click( | |
| start_document_rag_job, | |
| inputs=[ | |
| document_question, | |
| document_session_id, | |
| top_k_slider, | |
| max_loops_slider, | |
| resumable_requested_job_id, | |
| ], | |
| outputs=resumable_job_id, | |
| api_name="start_document_rag_job", | |
| api_visibility="public", | |
| queue=False, | |
| ) | |
| resumable_status_button.click( | |
| get_rag_job_status, | |
| inputs=resumable_job_id, | |
| outputs=resumable_job_status, | |
| api_name="get_rag_job_status", | |
| api_visibility="public", | |
| queue=False, | |
| ) | |
| resumable_forget_button.click( | |
| forget_rag_job, | |
| inputs=resumable_job_id, | |
| outputs=resumable_job_status, | |
| api_name="forget_rag_job", | |
| api_visibility="public", | |
| queue=False, | |
| ) | |
| question_box.submit( | |
| run_rag_agent, | |
| inputs, | |
| outputs, | |
| api_visibility="private", | |
| ) | |
| process_document_button.click( | |
| process_uploaded_document, | |
| inputs=[ | |
| uploaded_document, | |
| document_requested_session_id, | |
| ], | |
| outputs=[ | |
| document_session_id, | |
| document_status, | |
| vector_download, | |
| document_ready, | |
| ], | |
| api_name="process_uploaded_document", | |
| api_visibility="public", | |
| ) | |
| run_document_button.click( | |
| run_document_rag, | |
| inputs=[ | |
| document_question, | |
| document_session_id, | |
| top_k_slider, | |
| max_loops_slider, | |
| ], | |
| outputs=outputs, | |
| api_name="run_document_rag", | |
| api_visibility="public", | |
| ) | |
| validate_document_session_button.click( | |
| validate_document_session, | |
| inputs=document_session_id, | |
| outputs=[ | |
| document_session_valid, | |
| document_session_validation_message, | |
| ], | |
| api_name="validate_document_session", | |
| api_visibility="public", | |
| queue=False, | |
| ) | |
| clear_document_button.click( | |
| clear_document_session, | |
| inputs=document_session_id, | |
| outputs=document_status, | |
| api_name="clear_document_session", | |
| api_visibility="public", | |
| queue=False, | |
| ) | |
| clear_button.click( | |
| clear_interface, | |
| outputs=[ | |
| question_box, | |
| answer_output, | |
| process_output, | |
| sources_output, | |
| diagnostics_output, | |
| system_status, | |
| ], | |
| api_name="clear_interface", | |
| queue=False, | |
| ) | |
| status_button.click( | |
| check_system_status, | |
| outputs=system_status, | |
| api_name="check_system_status", | |
| queue=False, | |
| ) | |
| if __name__ == "__main__": | |
| demo.queue( | |
| default_concurrency_limit=DEFAULT_CONCURRENCY_LIMIT | |
| ).launch( | |
| debug=True, | |
| show_error=True, | |
| theme=gr.themes.Glass(primary_hue="blue"), | |
| css=CUSTOM_CSS, | |
| allowed_paths=[str(UPLOAD_SESSION_ROOT.resolve())], | |
| ssr_mode=False, | |
| ) | |