Update app.py

app.py

@@ -1,21 +1,20 @@
 # app.py
 # This is the updated main script. Copy-paste this over your existing app.py.
 # Changes:
-# - Removed old municipal loading code (e.g., MUNICIPAL_EMBEDDINGS_PATH, run_municipal_embedding_script_once, etc.) since the pre-built FAISS index is now downloaded from Hugging Face.
-# - Kept the combined RAG logic in route_model, which already merges CAP and municipal results.
-# - Updated prompt logic in rag_context to include municipal results seamlessly (no changes needed, as it's already combined).
-# - Jurisdiction filtering applies to combined results, checking state in citation/name.
-# - No changes to routing logic, as classify_prompt and route_model already handle task types appropriately.
-# - Retained preload_clusters for CAP clusters.
-# - Updated UI to match the provided image styling for desktop and mobile (responsive).
-# - Added support for .docx and .txt file uploads in addition to .pdf.
-# - Updated extract_text function to handle .pdf, .docx, .txt.
-# - Integrated file action dropdown into chat_interface logic.
-# - Added CSS for exact visual match, including dark blue theme, chat bubbles, and responsive design.
-# - Added header with VerdictAI logo using HTML (gavel emoji for simplicity; replace with image if needed).
-# - Removed image file type since not handled for legal analysis; focused on pdf, docx, txt.
-
-import gradio as gr
+# - Switched from Gradio to Flask for serving the custom HTML+CSS+JS frontend.
+# - Added API endpoint /api/chat for handling user inputs (prompt, jurisdiction, IRAC mode, web search toggle, file).
+# - Serves index.html as the root page (you'll need to add index.html to your repo with the provided HTML code).
+# - Integrated file handling in API (extracts text and appends to prompt if needed).
+# - Forced task_type to "irac" if IRAC mode is enabled; otherwise, uses classify_prompt.
+# - Added web_search toggle handling.
+# - Updated ask_gpt41_mini to use the fine-tuned model ft:gpt-4.1-mini-2025-04-14:w-jeffrey-scott-psc:verdictaitrain:BysFkyX4.
+# - If the task is document_creation, routes directly to the fine-tuned GPT model.
+# - Retained all other logic, including RAG (semantic_search for CAP + municipal_search for municipal; now hybrid with BM25 for municipal).
+# - Note: Add 'bm25s' to your requirements.txt for hybrid search (pip install bm25s).
+# - Note: The SaulLM endpoint is kept as-is (likely 7B; if you want 141B, update SAUL_ENDPOINT to a new HF cloud endpoint for SaulLM-141B).
+# - Note: For full chat history, the frontend JS handles appending messages client-side (stateless backend).
+
+import gradio as gr  # Retained if needed, but not used for UI anymore
 from openai import OpenAI
 import requests
 import os
@@ -44,6 +43,15 @@ from retrieval import *
 from prompt_builder import *
 from post_processing import *
 
+# Flask imports
+from flask import Flask, request, jsonify, send_from_directory
+from werkzeug.utils import secure_filename
+
+# BM25 for hybrid search (add 'bm25s' to requirements.txt)
+from bm25s import BM25
+
+app_flask = Flask(__name__)  # Renamed to avoid conflict with 'app' variable
+
 os.environ["HF_HOME"] = "/data/.huggingface"
 # Add or update this section in script.py
 # Ensure this is placed after imports but before any dataset loading or function definitions
@@ -195,7 +203,11 @@ def route_model(prompt, task_type, files=None, search_web=False, jurisdiction="K
     
     prompt = f"User prompt: {prompt}\n\n{rag_context}"
     
-    saul_response = ask_saul(prompt, task_type, jurisdiction)
+    if task_type == "document_creation":
+        # Route directly to fine-tuned GPT for document creation
+        saul_response = ask_gpt41_mini(prompt, jurisdiction)
+    else:
+        saul_response = ask_saul(prompt, task_type, jurisdiction)
     
     # Task-specific processing (existing code)
     saul_response = process_task_response(task_type, saul_response, prompt, jurisdiction)
@@ -237,26 +249,6 @@ def ask_saul(messages, task_type, jurisdiction):
         logger.error(f"SaulLM error: {str(e)}")
         return "SaulLM service unavailable. Using fallback response."
 
-def ask_gpt41_mini(prompt, jurisdiction):
-    try:
-        response = openai_client.chat.completions.create(
-            model="gpt-4",  # Placeholder, replace with fine-tuned model
-            messages=[
-                {"role": "system", "content": (
-                    f"You are a legal assistant drafting documents for {jurisdiction} jurisdiction. "
-                    "Always quote directly from retrieved case law. Use full case names and citations (e.g., 'Smith v. Jones, 123 S.W.3d 456 (Ky. 2005)'). "
-                    "Prioritize high quote density and include facts from those cases when applying them. Use IRAC structure. Do not paraphrase available holdings."
-                )},
-                {"role": "user", "content": prompt}
-            ],
-            temperature=0.3,
-            max_tokens=8192
-        )
-        return response.choices[0].message.content
-    except Exception as e:
-        logger.error(f"GPT-4.1 Mini error: {str(e)}")
-        return f"[GPT-4.1 Mini Error] {str(e)}"
-
 def ask_gpt4o(prompt):
     try:
         response = openai_client.chat.completions.create(
@@ -332,190 +324,68 @@ def classify_prompt(prompt):
         return "legal_strategy"
     return "general_qa"
 
-def chat_interface(prompt, files, history, search_web=False, jurisdiction="KY", action=None):
-    timestamp = datetime.now().strftime("%I:%M %p %m/%d/%Y")
-    task_type = classify_prompt(prompt)
+def summarize_document(file_text):
+    if file_text:
+        summary = ask_gpt4o(f"Summarize the following document: {file_text[:10000]}")  # Limit to avoid token limits
+        return f"Summary: {summary}"
+    return "No text extracted from file."
+
+def analyze_document(file_text):
+    if file_text:
+        analysis = ask_gpt4o(f"Analyze the following document for legal issues, risks, or key clauses: {file_text[:10000]}")
+        return f"Analysis: {analysis}"
+    return "No text extracted from file."
+
+def check_issues(file_text):
+    if file_text:
+        issues = ask_gpt4o(f"Check for red flags, unusual clauses, or potential issues in this legal document and highlight them: {file_text[:10000]}")
+        return f"Highlighted Issues: {issues}"
+    return "No text extracted from file."
+
+# Flask routes
+@app_flask.route('/')
+def index():
+    return send_from_directory('.', 'index.html')
+
+@app_flask.route('/api/chat', methods=['POST'])
+def api_chat():
+    prompt = request.form.get('prompt', '')
+    jurisdiction = request.form.get('jurisdiction', 'KY')
+    irac_mode = request.form.get('irac_mode', 'false') == 'true'
+    search_web = request.form.get('web_search', 'false') == 'true'
+    file = request.files.get('file')
+
     file_text = ""
-    if files:
-        file_text = extract_text_from_file(files[0]) if files else ""
-    if files and action:
-        if action == "Summarize":
-            response = summarize_document(files)
-        elif action == "Analyze":
-            response = analyze_document(files)
-        elif action == "Check Issues":
-            response = check_issues(files)
-        else:
-            prompt += "\nAttached file content: " + file_text[:10000]
-            response = route_model(prompt, task_type, files, search_web, jurisdiction)
-    elif files:
+    files = None
+    if file:
+        filename = secure_filename(file.filename)
+        temp_path = os.path.join('/tmp', filename)
+        file.save(temp_path)
+        file_text = extract_text_from_file(temp_path)
+        files = [temp_path]  # Pass as list for route_model
+        os.remove(temp_path)
+
+    task_type = classify_prompt(prompt)
+    if irac_mode:
+        task_type = "irac"
+
+    # Append file text to prompt if present
+    if file_text:
         if "summarize" in prompt.lower():
-            task_type = "document_analysis"
-            response = summarize_document(files)
+            response = summarize_document(file_text)
         elif "analyze" in prompt.lower():
-            task_type = "document_analysis"
-            response = analyze_document(files)
+            response = analyze_document(file_text)
         elif "check" in prompt.lower() or "issues" in prompt.lower() or "highlight" in prompt.lower():
-            task_type = "document_analysis"
-            response = check_issues(files)
+            response = check_issues(file_text)
         elif "generate" in prompt.lower() or "draft" in prompt.lower():
-            task_type = "document_creation"
-            response = ask_gpt41_mini(prompt + "\nAttached file content: " + file_text, jurisdiction)
+            response = ask_gpt41_mini(prompt + "\nAttached file content: " + file_text[:10000], jurisdiction)
         else:
             prompt += "\nAttached file content: " + file_text[:10000]
             response = route_model(prompt, task_type, files, search_web, jurisdiction)
     else:
         response = route_model(prompt, task_type, files, search_web, jurisdiction)
-    history.append((f"{prompt} <span style='color: #ECF0F1; font-size: 16px;'>[{timestamp}]</span>", 
-                    f"{response} <span style='color: #ECF0F1; font-size: 16px;'>[{timestamp}]</span>"))
-    return history, history
-
-def new_chat():
-    return [], []
-
-def summarize_document(files):
-    if files and isinstance(files, list) and files:
-        file = files[0]
-        text = extract_text_from_file(file)
-        if text:
-            summary = ask_gpt4o(f"Summarize the following document: {text[:10000]}")  # Limit to avoid token limits
-            return f"Summary: {summary}"
-        return "No text extracted from file."
-    return "Please upload a file to summarize."
-
-def analyze_document(files):
-    if files:
-        text = extract_text_from_file(files[0])
-        if text:
-            analysis = ask_gpt4o(f"Analyze the following document for legal issues, risks, or key clauses: {text[:10000]}")
-            return f"Analysis: {analysis}"
-        return "No text extracted from file."
-    return "No file uploaded for analysis."
-
-def check_issues(files):
-    if files:
-        text = extract_text_from_file(files[0])
-        if text:
-            issues = ask_gpt4o(f"Check for red flags, unusual clauses, or potential issues in this legal document and highlight them: {text[:10000]}")
-            return f"Highlighted Issues: {issues}"
-        return "No text extracted from file."
-    return "No file uploaded to check."
-
-def save_conversation(history):
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    content = "\n".join([f"User: {msg[0]}\nBot: {msg[1]}\n" for msg in history])
-    with open(f"conversation_{timestamp}.txt", "w") as f:
-        f.write(content)
-    return f"conversation_{timestamp}.txt"
-
-css = """
-body {
-    background-color: #2C3E50 !important;
-}
-.gradio-container {
-    background-color: #2C3E50 !important;
-    color: #ECF0F1 !important;
-}
-#chat-container {
-    background-color: #2C3E50;
-    height: 80vh;
-    overflow-y: auto;
-}
-.gr-chatbot {
-    background-color: #2C3E50 !important;
-}
-.gr-chatbot .message {
-    border-radius: 20px !important;
-    padding: 10px 15px !important;
-    max-width: 70% !important;
-    margin: 10px !important;
-    color: #ECF0F1 !important;
-}
-.gr-chatbot .message.user {
-    background-color: #34495E !important;
-    align-self: flex-end !important;
-    margin-left: auto !important;
-}
-.gr-chatbot .message.assistant {
-    background-color: #34495E !important;
-    align-self: flex-start !important;
-    margin-right: auto !important;
-}
-#chat-input {
-    background-color: #2C3E50 !important;
-    padding: 10px !important;
-    display: flex !important;
-    flex-wrap: wrap !important;
-}
-#user-input {
-    background-color: #34495E !important;
-    color: #ECF0F1 !important;
-    border: none !important;
-    flex-grow: 1 !important;
-    margin-right: 10px !important;
-}
-#send-btn {
-    background-color: #34495E !important;
-    color: #ECF0F1 !important;
-    border: none !important;
-    border-radius: 20px !important;
-}
-#file-upload-main {
-    background-color: #34495E !important;
-    color: #ECF0F1 !important;
-}
-#header {
-    background-color: #2C3E50 !important;
-    text-align: center !important;
-    padding: 20px !important;
-    font-size: 24px !important;
-    color: #AED6F1 !important;
-}
-@media (max-width: 768px) {
-    #chat-input {
-        flex-direction: column !important;
-    }
-    #chat-input > * {
-        margin-bottom: 10px !important;
-        width: 100% !important;
-    }
-    .gr-chatbot .message {
-        max-width: 90% !important;
-    }
-}
-"""
-
-theme = gr.themes.Base(
-    primary_hue="blue",
-    secondary_hue="blue",
-    neutral_hue="slate",
-).set(
-    body_text_color="#ECF0F1",
-    background_fill_primary="#2C3E50",
-    block_background_fill="#34495E",
-    input_background_fill="#34495E",
-    button_primary_background_fill="#34495E",
-    button_primary_text_color="#ECF0F1",
-)
-
-with gr.Blocks(css=css, theme=theme, title="VerdictAI - Legal Assistant") as app:
-    jurisdiction = gr.State("KY")
-    gr.HTML('<div id="header">🔨 VerdictAI</div>')  # Using hammer emoji as placeholder for gavel; replace with actual gavel icon if available
-    chatbot = gr.Chatbot(elem_id="chat-container", label="Chat")
-    with gr.Row(elem_id="chat-input"):
-        msg = gr.Textbox(
-            placeholder="Ask any legal question, request a draft document, upload a contract for analysis, or search for statutes and cases.\nExamples:\n‘Write a Kentucky will for a single parent with two children.’\n‘Summarize this operating agreement and flag any unusual clauses.’\n‘Find cases on constructive trust involving fraud.’\n‘What does KRS 411.182 mean for comparative fault?’\n‘IRAC analysis: A customer slips on an icy sidewalk outside a store.’",
-            elem_id="user-input"
-        )
-        file_upload = gr.File(file_count="multiple", file_types=[".pdf", ".docx", ".txt"], elem_id="file-upload-main", label="📎 Upload")
-        btn = gr.Button("Send", elem_id="send-btn")
-        google_search_btn = gr.Button("Google Search", elem_id="google-search-btn")
-        save_btn = gr.Button("Save Chat", elem_id="save-btn")
-        action_dropdown = gr.Dropdown(["Summarize", "Analyze", "Check Issues"], label="File Action")
-
-    btn.click(fn=chat_interface, inputs=[msg, file_upload, chatbot, gr.State(False), jurisdiction, action_dropdown], outputs=[chatbot, chatbot])
-    google_search_btn.click(fn=chat_interface, inputs=[msg, file_upload, chatbot, gr.State(True), jurisdiction, action_dropdown], outputs=[chatbot, chatbot])
-    save_btn.click(save_conversation, inputs=[chatbot], outputs=gr.File())
-
-    logger.info("Gradio app initialized successfully")
-    app.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)
+
+    return jsonify({'response': response})
+
+if __name__ == '__main__':
+    app_flask.run(host='0.0.0.0', port=7860)