Update App.py

App.py

@@ -1,31 +1,18 @@
-# script.py
-# This is the updated main script. Copy-paste this over your existing script.py.
-# Changes:
-# - Fixed any potential issues with prompt_lower by ensuring consistent use of prompt.lower() (though it was already correct).
-# - Split route_model responsibilities: Moved retrieval functions to retrieval.py, prompt building to prompt_builder.py, post-processing to post_processing.py.
-# - Imports: Add 'from retrieval import *', 'from prompt_builder import *', 'from post_processing import *' after your existing imports.
-# - For synchronous loads: hf_hub_download already caches, so subsequent calls are fast. To avoid first-time blocking, I've added a background preloader thread that pre-downloads a limited number of clusters (e.g., first 10) at startup. If you have many clusters, create a separate script (see below) to pre-download all offline.
-# - No need for a separate script unless you want to pre-download ALL clusters (which could be storage-intensive). If yes, see the optional pre_download_clusters.py below.
-# - In route_model, now calls the split functions.
-# - Passed necessary globals (e.g., cap_dataset, cap_id_to_index) to semantic_search.
-# - Ensured municipal_embeddings is loaded (assume it's global).
-# - Added threading for preloading some clusters at startup.
-
-import gradio as gr
-from openai import OpenAI
+# app.py
 import requests
 import os
 import logging
 from datetime import datetime
 import pdfplumber
-from googleapiclient.discovery import build
+from docx import Document
+from docx.shared import Pt, Inches
+from docx.enum.text import WD_ALIGN_PARAGRAPH
 import re
-from datasets import load_dataset, Dataset, load_from_disk
+from datasets import load_dataset, load_from_disk
 from sentence_transformers import SentenceTransformer
 import torch
 import numpy as np
 import shutil
-import pyarrow.parquet as pq
 from huggingface_hub import hf_hub_download
 import pickle
 import faiss
@@ -33,340 +20,221 @@ import threading
 import subprocess
 from task_processing import process_task_response
 from gpt_helpers import ask_gpt41_mini
+from retrieval import retrieve_context
+from prompt_builder import build_grok_prompt, build_editor_prompt
+from flask import Flask, request, jsonify, send_from_directory, send_file, Response, stream_with_context
+from werkzeug.utils import secure_filename
+from rank_bm25 import BM25Okapi
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+import json  # For safer JSON parsing if needed
 
-# New imports for split modules
-from retrieval import *
-from prompt_builder import *
-from post_processing import *
+app = Flask(__name__)  # Renamed from app_flask to app for HF Spaces compatibility
 
 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
 
-from huggingface_hub import login
+# Logging setup
+logger = logging.getLogger("app")
+logging.basicConfig(level=logging.INFO)
+logger.info("✅ Logging initialized. Starting app setup.")
+print("App setup starting...")  # Fallback print for early debug
 
-# Load HF token for SaulLM endpoint and gated repos
+# Hugging Face authentication
+from huggingface_hub import login
 hf_token = os.environ.get("HF_TOKEN", "")
-if not hf_token:
-    logger.warning("HF_TOKEN not set; SaulLM endpoint may require authentication and gated repos may not be accessible.")
-
-# Authenticate for gated Hugging Face repos (e.g., for centroids download)
 if hf_token:
     login(hf_token)
-    logger.info("Authenticated with Hugging Face token for gated repos.")
+    logger.info("✅ Authenticated with Hugging Face token for gated repos.")
 else:
-    logger.warning("No HF_TOKEN; may fail to access gated repos like Caselaw_Access_Project_embeddings.")
+    logger.warning("HF_TOKEN not set; gated repos may not be accessible.")
 
 # Check environment variables
-try:
-    OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "Missing")
-    GOOGLE_SEARCH_API = os.environ.get("GOOGLE_SEARCH_API", "Missing")  # This is now treated as CSE ID (cx)
-    GOOGLE_CUSTOM_SEARCH_API_KEY = os.environ.get("GOOGLE_CUSTOM_SEARCH_API_KEY", "Missing")  # New: API key (developerKey)
-    if OPENAI_API_KEY == "Missing" or GOOGLE_CUSTOM_SEARCH_API_KEY == "Missing" or GOOGLE_SEARCH_API == "Missing":
-        raise KeyError("API keys not set")
-    logger.info(f"OpenAI API Key starts with: {OPENAI_API_KEY[:10]}...")
-    logger.info("API keys loaded successfully")
-except KeyError as e:
-    logger.error(f"Missing environment variable: {str(e)}")
-    raise EnvironmentError(f"Required secrets OPENAI_API_KEY, GOOGLE_CUSTOM_SEARCH_API_KEY, and GOOGLE_SEARCH_API must be set in Hugging Face Space Secrets")
-def download_file_if_not_exists(url, save_path):
-    if not os.path.exists(save_path):
-        try:
-            response = requests.get(url)
-            response.raise_for_status()  # Raise error if not 200
-            with open(save_path, 'wb') as f:
-                f.write(response.content)
-            logger.info(f"Downloaded and saved file to {save_path}")
-        except Exception as e:
-            logger.error(f"Failed to download from {url}: {str(e)}")
-
-# Download the centroids file if not present
-centroid_url = "https://huggingface.co/datasets/laion/Caselaw_Access_Project_embeddings/blob/main/TeraflopAI___Caselaw_Access_Project_centroids.parquet"
-centroid_path = "TeraflopAI___Caselaw_Access_Project_centroids.parquet"
-download_file_if_not_exists(centroid_url, centroid_path)
-# Load HF token for SaulLM endpoint
-hf_token = os.environ.get("HF_TOKEN", "")
-if not hf_token:
-    logger.warning("HF_TOKEN not set; SaulLM endpoint may require authentication")
-
-import requests
-
-
-# Initialize OpenAI client
-openai_client = OpenAI(api_key=OPENAI_API_KEY)
-
-# SaulLM endpoint
-SAUL_ENDPOINT = "https://l4tuv4j9bu616t5x.us-east-1.aws.endpoints.huggingface.cloud"
-
-# Persistent storage path for dataset
-LOCAL_PATH = "/data/cap_dataset"
-dataset_info_path = os.path.join(LOCAL_PATH, 'dataset_info.json')
-if os.path.exists(dataset_info_path):
-    cap_dataset = load_from_disk(LOCAL_PATH)
+OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "Missing")
+GOOGLE_SEARCH_API = os.environ.get("GOOGLE_SEARCH_API", "Missing")  # CSE ID
+GOOGLE_CUSTOM_SEARCH_API_KEY = os.environ.get("GOOGLE_CUSTOM_SEARCH_API_KEY", "Missing")  # API key
+COURT_LISTENER_API_KEY = os.environ.get("Court_Listener_API", "Missing")  # Updated to match HF secret name
+if OPENAI_API_KEY == "Missing":
+    logger.warning("OPENAI_API_KEY not set; OpenAI features will fail.")
+if GOOGLE_CUSTOM_SEARCH_API_KEY == "Missing" or GOOGLE_SEARCH_API == "Missing":
+    logger.warning("Google Search keys not set; search features will fail.")
+if COURT_LISTENER_API_KEY == "Missing":
+    logger.warning("Court_Listener_API not set; CourtListener features will fail.")
+logger.info("✅ API keys checked (with warnings if missing).")
+
+# Initialize OpenAI client (only if key present)
+openai_client = None
+if OPENAI_API_KEY != "Missing":
+    from openai import OpenAI
+    openai_client = OpenAI(api_key=OPENAI_API_KEY)
+    logger.info("✅ OpenAI client initialized.")
 else:
-    try:
-        cap_dataset = load_dataset("TeraflopAI/Caselaw-Access-Project", split="train")
-        cap_dataset.save_to_disk(LOCAL_PATH)
-    except Exception as e:
-        logger.error(f"Dataset download/save failed: {str(e)}")
-        if os.path.exists(LOCAL_PATH):
-            shutil.rmtree(LOCAL_PATH)  # Clean up partial save
-        raise
-
-# Persistent storage paths for municipal
-MUNICIPAL_EMBEDDINGS_PATH = "/data/municipal_embeddings"
-MUNICIPAL_HTML_PATH = "/data/municipal_html"
-MUNICIPAL_CITATION_PATH = "/data/municipal_citation"
-MUNICIPAL_FAISS_INDEX_PATH = "/data/municipal_faiss.index"
-MUNICIPAL_CID_TO_HTML_PATH = "/data/cid_to_html.pkl"
-MUNICIPAL_CID_TO_CITATION_PATH = "/data/cid_to_citation.pkl"
-
-# Only trigger once to avoid re-downloading
-def run_municipal_embedding_script_once():
-    marker_path = "/data/municipal_embeddings_done.txt"
-    if not os.path.exists(marker_path):
-        try:
-            subprocess.run(["python", "prepare_municipal_embeddings.py"], check=True)
-            with open(marker_path, "w") as f:
-                f.write("done")
-            logger.info("✅ Municipal embedding preparation complete.")
-        except Exception as e:
-            logger.error(f"❌ Error running prepare_municipal_embeddings.py: {e}")
-
-# Only trigger once to avoid re-building
-def run_build_municipal_faiss_once():
-    marker_path = "/data/municipal_faiss_done.txt"
-    if not os.path.exists(marker_path):
-        try:
-            subprocess.run(["python", "build_municipal_faiss.py"], check=True)
-            with open(marker_path, "w") as f:
-                f.write("done")
-            logger.info("✅ Municipal FAISS index build complete.")
-        except Exception as e:
-            logger.error(f"❌ Error running build_municipal_faiss.py: {e}")
-
-# Launch in background AFTER app boots
-# Run preparation scripts synchronously if not done
-run_municipal_embedding_script_once()
-run_build_municipal_faiss_once()
-
-# Load municipal embeddings dataset
-if os.path.exists(MUNICIPAL_EMBEDDINGS_PATH):
-    municipal_embeddings = load_from_disk(MUNICIPAL_EMBEDDINGS_PATH)
-else:
-    logger.error("Municipal embeddings not found. Ensure prepare_municipal_embeddings.py ran successfully.")
-    municipal_embeddings = None  # Fallback or error handling
-
-# Load municipal html dataset
-if os.path.exists(MUNICIPAL_HTML_PATH):
-    municipal_html = load_from_disk(MUNICIPAL_HTML_PATH)
-else:
-    logger.error("Municipal html not found. Ensure prepare_municipal_embeddings.py ran successfully.")
-    municipal_html = None
-
-# Load municipal citation dataset
-if os.path.exists(MUNICIPAL_CITATION_PATH):
-    municipal_citation = load_from_disk(MUNICIPAL_CITATION_PATH)
-else:
-    logger.error("Municipal citation not found. Ensure prepare_municipal_embeddings.py ran successfully.")
-    municipal_citation = None
-
-# Precompute CID to index mapping for CAP dataset
-cap_id_to_index = {doc['cid']: i for i, doc in enumerate(cap_dataset) if 'cid' in doc}
-
-# Preload some clusters in background (e.g., clusters 0-9)
-def preload_clusters():
-    for cluster_id in range(10):  # Adjust range as needed
-        try:
-            load_cluster_vectors(cluster_id, model="gte-large")
-            logger.info(f"Preloaded cluster {cluster_id}")
-        except Exception as e:
-            logger.error(f"Preload failed for cluster {cluster_id}: {e}")
+    logger.warning("Skipping OpenAI client init due to missing key.")
+
+# Grok API setup
+GROK_API_URL = "https://api.x.ai/v1/chat/completions"
+GROK_API_TOKEN = "xai-fr0jVd7v8jiuxAQko2rpx1ft7DIK0iQkHQTk0RSFubXxdKm6AUgho4hJnlJ2OINlh82AYJ4GELGqLMSb"  # From user
+logger.info("✅ Grok API endpoint and token set.")
+
+# Global session for retries
+session = requests.Session()
+retries = Retry(total=3, backoff_factor=1, status_forcelist=[422, 503, 504])
+session.mount('https://', HTTPAdapter(max_retries=retries))
+
+# Lazy-load CAP dataset to avoid startup issues
+def get_cap_dataset():
+    if not hasattr(get_cap_dataset, 'dataset') or get_cap_dataset.dataset is None:
+        from datasets import load_from_disk  # Lazy import
+        LOCAL_PATH = "/data/cap_dataset"
+        if os.path.exists(os.path.join(LOCAL_PATH, 'dataset_info.json')):
+            try:
+                get_cap_dataset.dataset = load_from_disk(LOCAL_PATH)
+                logger.info("✅ Lazy-loaded CAP dataset from /data/cap_dataset.")
+            except Exception as e:
+                logger.error(f"Failed to load CAP dataset: {str(e)}")
+                get_cap_dataset.dataset = None
+        else:
+            logger.error("CAP dataset not found at /data/cap_dataset. Ensure it’s preloaded.")
+            get_cap_dataset.dataset = None
+    return get_cap_dataset.dataset
+
+get_cap_dataset.dataset = None
+logger.info("✅ CAP dataset lazy-loader defined.")
+
+# Lazy-compute CID to index mapping for CAP dataset
+def get_cap_id_to_index():
+    if not hasattr(get_cap_id_to_index, 'index') or get_cap_id_to_index.index is None:
+        import pickle  # Lazy import if needed elsewhere, but here it's for logging only
+        cap_dataset = get_cap_dataset()
+        if cap_dataset is not None:
+            get_cap_id_to_index.index = {doc['cid']: i for i, doc in enumerate(cap_dataset) if 'cid' in doc}
+            logger.info("✅ Precomputed CAP CID to index mapping.")
+        else:
+            get_cap_id_to_index.index = {}
+            logger.error("CAP dataset not available for index mapping.")
+    return get_cap_id_to_index.index
 
-threading.Thread(target=preload_clusters).start()
+get_cap_id_to_index.index = None
+logger.info("✅ CAP ID-to-index lazy-loader defined.")
 
 # State dictionary for jurisdiction
 STATES = {
-    "AL": "Alabama",
-    "AK": "Alaska",
-    "AZ": "Arizona",
-    "AR": "Arkansas",
-    "CA": "California",
-    "CO": "Colorado",
-    "CT": "Connecticut",
-    "DE": "Delaware",
-    "FL": "Florida",
-    "GA": "Georgia",
-    "HI": "Hawaii",
-    "ID": "Idaho",
-    "IL": "Illinois",
-    "IN": "Indiana",
-    "IA": "Iowa",
-    "KS": "Kansas",
-    "KY": "Kentucky",
-    "LA": "Louisiana",
-    "ME": "Maine",
-    "MD": "Maryland",
-    "MA": "Massachusetts",
-    "MI": "Michigan",
-    "MN": "Minnesota",
-    "MS": "Mississippi",
-    "MO": "Missouri",
-    "MT": "Montana",
-    "NE": "Nebraska",
-    "NV": "Nevada",
-    "NH": "New Hampshire",
-    "NJ": "New Jersey",
-    "NM": "New Mexico",
-    "NY": "New York",
-    "NC": "North Carolina",
-    "ND": "North Dakota",
-    "OH": "Ohio",
-    "OK": "Oklahoma",
-    "OR": "Oregon",
-    "PA": "Pennsylvania",
-    "RI": "Rhode Island",
-    "SC": "South Carolina",
-    "SD": "South Dakota",
-    "TN": "Tennessee",
-    "TX": "Texas",
-    "UT": "Utah",
-    "VT": "Vermont",
-    "VA": "Virginia",
-    "WA": "Washington",
-    "WV": "West Virginia",
-    "WI": "Wisconsin",
-    "WY": "Wyoming",
-    "Federal": "Federal",
-    "All States": "All States",
-    "Other": "Other States"
+    "AL": "Alabama", "AK": "Alaska", "AZ": "Arizona", "AR": "Arkansas", "CA": "California",
+    "CO": "Colorado", "CT": "Connecticut", "DE": "Delaware", "FL": "Florida", "GA": "Georgia",
+    "HI": "Hawaii", "ID": "Idaho", "IL": "Illinois", "IN": "Indiana", "IA": "Iowa",
+    "KS": "Kansas", "KY": "Kentucky", "LA": "Louisiana", "ME": "Maine", "MD": "Maryland",
+    "MA": "Massachusetts", "MI": "Michigan", "MN": "Minnesota", "MS": "Mississippi", "MO": "Missouri",
+    "MT": "Montana", "NE": "Nebraska", "NV": "Nevada", "NH": "New Hampshire", "NJ": "New Jersey",
+    "NM": "New Mexico", "NY": "New York", "NC": "North Carolina", "ND": "North Dakota", "OH": "Ohio",
+    "OK": "Oklahoma", "OR": "Oregon", "PA": "Pennsylvania", "RI": "Rhode Island", "SC": "South Carolina",
+    "SD": "South Dakota", "TN": "Tennessee", "TX": "Texas", "UT": "Utah", "VT": "Vermont",
+    "VA": "Virginia", "WA": "Washington", "WV": "West Virginia", "WI": "Wisconsin", "WY": "Wyoming",
+    "Federal": "Federal", "All States": "All States", "Other": "Other States"
 }
+logger.info("✅ States dictionary loaded.")
 
-def route_model(prompt, task_type, files=None, search_web=False, jurisdiction="KY"):
-    logger.info(f"Routing prompt: {prompt}, Task: {task_type}, Web Search: {search_web}, Jurisdiction: {jurisdiction}")
-    
-    rag_context = ""
-    if task_type in ["case_law", "irac", "statute"]:
-        cap_results = semantic_search(prompt, top_k=5)
-        municipal_results = municipal_search(prompt, top_k=5)
-        combined_results = cap_results + municipal_results
-        
-        # Filter by jurisdiction if specified (e.g., "KY" for Kentucky)
-        if jurisdiction and jurisdiction != "All States":
-            state_name = STATES.get(jurisdiction, "")
-            state_code = jurisdiction  # e.g., "KY"
-            combined_results = [r for r in combined_results if state_code in r['citation'] or state_name in r['citation'] or state_code in r['name'] or state_name in r['name']]
-        
-        if combined_results:
-            rag_context = "Retrieved legal authorities (case law and statutes):\n" + "\n".join([f"{i+1}. [{auth.get('source', 'Unknown')}] {auth['name']}, {auth['citation']}: \"{auth['snippet']}\"" for i, auth in enumerate(combined_results)])
-    
-    prompt = f"User prompt: {prompt}\n\n{rag_context}"
-    
-    saul_response = ask_saul(prompt, task_type, jurisdiction)
-    
-    # Task-specific processing (existing code)
-    saul_response = process_task_response(task_type, saul_response, prompt, jurisdiction)
-    
-    if search_web:
-        web_data = google_search(prompt)
-        saul_response = f"Google Search results: {web_data}\n{saul_response}"
-
-    editor_prompt = build_editor_prompt(prompt, task_type, jurisdiction, saul_response, rag_context)
-    
-    final_response = ask_gpt4o(editor_prompt)
-    
-    final_response = ground_statutes(final_response, jurisdiction)
-    
-    return final_response
-
-def ask_saul(messages, task_type, jurisdiction):
+# Verdict Ai api call function (updated for streaming)
+def ask_grok(messages, stream=False):
     try:
-        headers = {"Authorization": f"Bearer {hf_token}"} if hf_token else {}
+        headers = {
+            "Accept": "application/json",
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {GROK_API_TOKEN}"
+        }
         payload = {
             "messages": messages,
-            "parameters": {
-                "max_length": 512,  
-                "temperature": 0.3
+            "model": "grok-4-0709",
+            "stream": stream,
+            "temperature": 0.1,
+            "max_tokens": 131072,  # High value for long responses
+            "search_parameters": {
+                "mode": "on"
             }
         }
-        logger.info(f"SaulLM payload: messages length={len(messages)}, max_length={payload['parameters']['max_length']}")
-        response = requests.post(SAUL_ENDPOINT, headers=headers, json=payload)
+        logger.info(f"Grok payload: {payload}")
+        response = requests.post(GROK_API_URL, headers=headers, json=payload, stream=stream)
+        logger.info(f"Grok response status: {response.status_code}")
         response.raise_for_status()
-        result = response.json()
-        if isinstance(result, dict) and "choices" in result:
-            return result["choices"][0].get("message", {}).get("content", "[No response from SaulLM]")
-        elif isinstance(result, list) and result:
-            return result[0].get("generated_text", "[No response from SaulLM]")
+        if stream:
+            def stream_gen():
+                logger.info("Starting Grok stream...")
+                for raw_chunk in response.iter_lines():
+                    chunk = raw_chunk.decode("utf-8").strip()
+                    if not chunk:
+                        continue  # Skip empty lines
+                    chunk_data = chunk.replace("data: ", "")
+                    logger.info(f"Raw chunk: {chunk_data}")
+                    if chunk_data == "[DONE]":
+                        yield "data: [DONE]\n\n"
+                        break
+                    try:
+                        result = json.loads(chunk_data)
+                        delta = result.get("choices", [{}])[0].get("delta", {})
+                        content = delta.get("content", "")
+                        if content:
+                            yield f'data: {{"chunk": {json.dumps(content)}}}\n\n'
+                    except Exception as e:
+                        logger.warning(f"Grok JSON parse error: {e} | chunk_data: {chunk_data}")
+                        yield f'data: {{"chunk": "[Unrecognized Grok output]"}}\n\n'
+                logger.info("Stream ended.")
+            return stream_gen()
         else:
-            return result.get("generated_text", "[No response from SaulLM]")
-
+            result = response.json()
+            logger.info(f"Grok non-stream result: {result}")
+            if "choices" in result and result["choices"] and "message" in result["choices"][0] and "content" in result["choices"][0]["message"]:
+                content = result["choices"][0]["message"]["content"]
+                if len(content) > 65536:
+                    content = content[:65536] + "... [Truncated]"
+                return content.strip()
+            return "[No response]"
+    except requests.exceptions.HTTPError as http_err:
+        logger.error(f"Grok HTTP error: {http_err}, Response: {response.text if 'response' in locals() else 'N/A'}")
+        if stream:
+            def error_gen():
+                yield f'data: {{"error": "Grok API error: {str(http_err)}"}}\n\n'
+                yield "data: [DONE]\n\n"
+            return error_gen()
+        return "[Grok Error] " + str(http_err)
     except Exception as e:
-        logger.error(f"SaulLM error: {str(e)}")
-        return "SaulLM service unavailable. Using fallback response."
-
-def ask_gpt41_mini(prompt, jurisdiction):
+        logger.error(f"Grok general error: {type(e).__name__}: {str(e)}")
+        if stream:
+            def error_gen():
+                yield f'data: {{"error": "{str(e)}"}}\n\n'
+                yield "data: [DONE]\n\n"
+            return error_gen()
+        return "[No response]"
+
+def extract_text_from_file(file_path):
     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(
-            model="gpt-4o",
-            messages=[
-                {"role": "system", "content": (
-                    "You are the final editor for a legal research assistant. Polish and organize the output into clear IRAC format. "
-                    "Ensure high quote density from retrieved authorities and include relevant facts from the cited cases. "
-                    "Maintain accurate citations. Do not paraphrase legal holdings when direct quotes are available."
-                )},
-                {"role": "user", "content": prompt}
-            ],
-            temperature=0.3,
-            max_tokens=16384
-        )
-        return response.choices[0].message.content
-    except Exception as e:
-        logger.error(f"GPT-4o error: {str(e)}")
-        return f"[GPT-4o Error] {str(e)}"
-
-def extract_text_from_pdf(file_path):
-    try:
-        with pdfplumber.open(file_path) as pdf:
-            text = ""
-            for page in pdf.pages:
-                text += page.extract_text() or ""
-            logger.info(f"Extracted text length: {len(text)}")
-            return text
+        ext = os.path.splitext(file_path)[1].lower()
+        text = ""
+        if ext == '.pdf':
+            with pdfplumber.open(file_path) as pdf:
+                text = "\n".join([page.extract_text() or "" for page in pdf.pages])
+        elif ext == '.docx':
+            doc = Document(file_path)
+            text = "\n".join([para.text for para in doc.paragraphs])
+        elif ext == '.txt':
+            with open(file_path, 'r', encoding='utf-8') as f:
+                text = f.read()
+        else:
+            text = f"Non-text file uploaded: {os.path.basename(file_path)}. Analyze if image or other."
+        logger.info(f"Extracted text length: {len(text)} from {ext} file")
+        return text
     except Exception as e:
-        logger.error(f"PDF extraction error: {str(e)}")
+        logger.error(f"File extraction error: {str(e)}")
         return ""
 
 def classify_prompt(prompt):
     prompt_lower = prompt.lower()
     if "summarize" in prompt_lower:
         return "document_analysis"  # Treat summarize as analysis for routing
-    if any(k in prompt_lower for k in ["irac", "issue", "rule", "analysis", "conclusion"]):
+    if any(k in prompt_lower for k in ["irac", "issue", "rule", "analysis", "conclusion", "brief", "memorandum", "memo"]):
         return "irac"
     elif any(k in prompt_lower for k in ["case", "precedent", "law"]):
         return "case_law"
     elif any(k in prompt_lower for k in ["statute", "krs"]):
         return "statute"
-    elif any(k in prompt_lower for k in ["draft", "write", "generate", "petition", "letter", "contract"]):
+    elif any(k in prompt_lower for k in ["draft", "write", "generate", "petition", "letter", "contract", "title opinion"]):
         return "document_creation"
     elif any(k in prompt_lower for k in ["review", "summarize", "clause", "red flags"]):
         return "document_analysis"
@@ -392,104 +260,321 @@ def classify_prompt(prompt):
         return "legal_strategy"
     return "general_qa"
 
-def chat_interface(prompt, files, history, search_web=False, jurisdiction="KY"):
-    timestamp = datetime.now().strftime("%I:%M %p %m/%d/%Y")
-    task_type = classify_prompt(prompt)
-    if files:
-        file_text = extract_text_from_pdf(files[0]) if files else ""
-        if "summarize" in prompt.lower():
-            task_type = "document_analysis"
-            response = summarize_document(files)
-        elif "analyze" in prompt.lower():
-            task_type = "document_analysis"
-            response = analyze_document(files)
-        elif "check" in prompt.lower() or "issues" in prompt.lower():
-            task_type = "document_analysis"
-            response = check_issues(files)
-        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)
+def create_legal_docx(content, jurisdiction, filename):
+    doc = Document()
+    # Set margins and font
+    sections = doc.sections
+    for section in sections:
+        section.top_margin = Inches(1)
+        section.bottom_margin = Inches(1)
+        section.left_margin = Inches(1)
+        section.right_margin = Inches(1)
+    # Case Caption (example placeholder)
+    caption = doc.add_paragraph()
+    caption.alignment = WD_ALIGN_PARAGRAPH.CENTER
+    run = caption.add_run("IN THE [COURT NAME] OF [JURISDICTION]\n")
+    run.bold = True
+    run.font.size = Pt(12)
+    caption.add_run("[Plaintiff] v. [Defendant]\nCase No: [Number]")
+    # Add content (assume content has sections marked with # for headings)
+    lines = content.split('\n')
+    for line in lines:
+        if line.startswith('# '):
+            heading = doc.add_heading(line[2:], level=1)
+            heading.alignment = WD_ALIGN_PARAGRAPH.CENTER
+        elif line.startswith('## '):
+            doc.add_heading(line[3:], level=2)
         else:
-            prompt += "\nAttached file content: " + file_text[:10000]
-            response = route_model(prompt, task_type, files, search_web, jurisdiction)
+            p = doc.add_paragraph(line)
+            p.alignment = WD_ALIGN_PARAGRAPH.JUSTIFY
+    # Signature Block
+    doc.add_paragraph("\nRespectfully submitted,")
+    sig = doc.add_paragraph("[Attorney Name]\n[Bar Number]\n[Firm]\n[Address]\n[Phone]\n[Email]")
+    sig.alignment = WD_ALIGN_PARAGRAPH.LEFT
+    # Certificate of Service
+    doc.add_heading("CERTIFICATE OF SERVICE", level=1)
+    doc.add_paragraph("I hereby certify that a true and correct copy of the foregoing was served on [date] via [method] to:\n[Recipient]")
+    # Notary Acknowledgement (if applicable)
+    doc.add_heading("NOTARY ACKNOWLEDGEMENT", level=1)
+    doc.add_paragraph("[State/County]\nSubscribed and sworn to before me this [date] by [name].\n\nNotary Public")
+    doc.save(filename)
+    return filename
+
+def route_model(messages, task_type, files=None, search_web=False, jurisdiction="KY"):
+    logger.info(f"Routing messages, Task: {task_type}, Web Search: {search_web}, Jurisdiction: {jurisdiction}")
+    rag_context = ""
+    prompt = messages[-1]['content']  # Use last user message as prompt for classification etc.
+    if task_type in ["case_law", "irac", "statute"]:  # Skip RAG for document_creation/summaries
+        cap_dataset = get_cap_dataset()
+        if cap_dataset is not None:
+            combined_results = retrieve_context(prompt, task_type, jurisdiction)
+            # Filter by jurisdiction if specified
+            if jurisdiction and jurisdiction != "All States":
+                state_name = STATES.get(jurisdiction, "").lower()
+                state_code = jurisdiction.lower()
+                variants = [state_code, state_name, f"{state_code}.", state_name.replace(" ", "")]
+                combined_results = [r for r in combined_results if any(v in (r.get('citation', '') + r.get('name', '') + r.get('snippet', '')).lower() for v in variants)]
+            if combined_results:
+                rag_context = "Retrieved legal authorities (case law and statutes):\n" + "\n".join(
+                    [f"{i+1}. [{auth.get('source', 'Unknown')}] {auth['name']}, {auth['citation']}: \"{auth['snippet']}\"" for i, auth in enumerate(combined_results)]
+                )
+                messages[-1]['content'] = f"{prompt}\n\n{rag_context}"
+    if task_type == "document_creation":
+        # Reset messages to only current prompt to avoid history accumulation
+        prompt = messages[-1]['content']
+        draft_messages = [{'role': 'user', 'content': prompt}]
+        # Route directly to fine-tuned GPT for document creation
+        gpt_response = ask_gpt41_mini(prompt, jurisdiction)  # Adjust to use full messages if gpt_helpers supports
+        logger.info(f"GPT-4.1-mini response length: {len(gpt_response)} | Content snippet: {gpt_response[:200]}...")
+        if not gpt_response.strip():
+            logger.warning("Empty response from GPT-4.1-mini; possible content filtering.")
+            yield f'data: {{"error": "Empty draft from GPT-4.1-mini - prompt may be filtered. Try rephrasing."}}\n\n'
+            yield "data: [DONE]\n\n"
+            return
+        # Truncate if too long to prevent token issues
+        MAX_GPT_LEN = 20000
+        if len(gpt_response) > MAX_GPT_LEN:
+            gpt_response = gpt_response[:MAX_GPT_LEN] + "\n[Truncated: GPT response too long; refining may be needed.]"
+            logger.warning(f"Truncated GPT response to {MAX_GPT_LEN} chars.")
+        editor_messages = draft_messages + [{'role': 'assistant', 'content': gpt_response}]
+        editor_prompt = build_editor_prompt(prompt, task_type, jurisdiction, gpt_response, rag_context)  # But to make contextual, perhaps use full
+        editor_messages.append({'role': 'user', 'content': editor_prompt})  # Or append
+        # Use non-stream for Grok to avoid streaming issues
+        try:
+            full_grok_response = ask_grok(editor_messages, stream=False)  # CHANGED: Non-stream for reliability
+            logger.info(f"Grok polish response length: {len(full_grok_response)} | Snippet: {full_grok_response[:200]}...")
+            if not full_grok_response.strip():
+                logger.warning("Empty response from Grok; using GPT draft.")
+                full_response = gpt_response
+            else:
+                full_response = full_grok_response
+        except Exception as e:
+            logger.error(f"Grok non-stream error: {str(e)}. Using GPT draft.")
+            full_response = gpt_response
+        # Yield as faux stream chunks
+        chunks = [full_response[i:i+200] for i in range(0, len(full_response), 200)]  # Split for streaming feel
+        for part in chunks:
+            yield f'data: {{"chunk": {json.dumps(part)}}}\n\n'  # Use json.dumps for safe escaping
+        # Create doc and send download URL
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        filename = f"/tmp/legal_doc_{timestamp}.docx"
+        create_legal_docx(full_response, jurisdiction, filename)
+        yield f'data: {{"download_url": "/download/legal_doc_{timestamp}.docx"}}\n\n'
+        yield "data: [DONE]\n\n"
+        return
     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
+        try:
+            # Build system prompt contextual
+            system_content = build_grok_prompt(prompt, task_type, jurisdiction, rag_context)  # But since messages have it, prepend if not
+            system_content += "\nStick strictly to the provided retrieved context for your response. Do not add information, cases, or statutes not explicitly in the context to avoid hallucinations. If context is insufficient, state so clearly."
+            if 'CourtListener' in rag_context:
+                system_content += "\nPrioritize CourtListener results for accuracy: Quote key snippets, cite cases, and polish into a structured response (e.g., IRAC format for analysis tasks)."
+            if messages[0]['role'] != 'system':
+                messages = [{'role': 'system', 'content': system_content}] + messages
+            stream_grok = ask_grok(messages, stream=True)
+        except Exception as e:
+            logger.error(f"Grok failed: {e}. Falling back to GPT-4o.")
+            grok_response = ask_gpt4o(messages[-1]['content'])  # Fallback, adjust to full if possible
+            yield f'data: {{"chunk": "{grok_response}"}}\n\n'
+            yield "data: [DONE]\n\n"
+            return
+        # Task-specific processing
+        # For streaming, skip or adapt; here, stream raw
+        for chunk in stream_grok:
+            yield chunk
+        yield "data: [DONE]\n\n"
 
-def new_chat():
-    return [], []
+def ask_gpt4o(prompt):
+    try:
+        irac_system = "If the task involves legal analysis, polish and organize the output into clear IRAC format. Otherwise, organize appropriately without IRAC."
+        response = openai_client.chat.completions.create(
+            model="gpt-4o",
+            messages=[
+                {
+                    "role": "system",
+                    "content": (
+                        f"You are the final editor for a legal research assistant. {irac_system} "
+                        "Ensure high quote density from retrieved authorities and include relevant facts from the cited cases. "
+                        "Maintain accurate citations. Do not paraphrase legal holdings when direct quotes are available. "
+                        "Do not cite or reference any case law, statutes, or authorities that are not explicitly provided in the retrieved context or user input."
+                    )
+                },
+                {"role": "user", "content": prompt}
+            ],
+            temperature=0.3,
+            max_tokens=65536
+        )
+        return response.choices[0].message.content
+    except Exception as e:
+        logger.error(f"GPT-4o error: {str(e)}")
+        return f"[GPT-4o Error] {str(e)}"
 
 def summarize_document(files):
-    if files and isinstance(files, list) and files:
-        file = files[0]
-        text = extract_text_from_pdf(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 PDF."
-    return "Please upload a file to summarize."
+    def gen():
+        if files and isinstance(files, list) and files:
+            texts = [extract_text_from_file(f) for f in files]
+            text = "\n".join(texts)
+            if text:
+                summary = ask_grok([{"role": "user", "content": f"Summarize the following document(s): {text[:10000]}"}], stream=False)  # Explicitly non-stream
+                full_response = f"Summary: {summary}"
+                chunks = [full_response[i:i+200] for i in range(0, len(full_response), 200)]  # Split for streaming feel
+                for part in chunks:
+                    yield f'data: {{"chunk": {json.dumps(part)}}}\n\n'
+                yield "data: [DONE]\n\n"
+            else:
+                yield f'data: {{"chunk": "No text extracted from file."}}\n\n'
+                yield "data: [DONE]\n\n"
+        else:
+            yield f'data: {{"chunk": "Please upload a file to summarize."}}\n\n'
+            yield "data: [DONE]\n\n"
+    return gen
 
 def analyze_document(files):
-    if files:
-        text = extract_text_from_pdf(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 PDF."
-    return "No file uploaded for analysis."
+    def gen():
+        if files:
+            texts = [extract_text_from_file(f) for f in files]
+            text = "\n".join(texts)
+            if text:
+                analysis = ask_grok([{"role": "user", "content": f"Analyze the following document(s) for legal issues, risks, or key clauses: {text[:10000]}"}], stream=False)  # Explicitly non-stream
+                full_response = f"Analysis: {analysis}"
+                chunks = [full_response[i:i+200] for i in range(0, len(full_response), 200)]
+                for part in chunks:
+                    yield f'data: {{"chunk": {json.dumps(part)}}}\n\n'
+                yield "data: [DONE]\n\n"
+            else:
+                yield f'data: {{"chunk": "No text extracted from file."}}\n\n'
+                yield "data: [DONE]\n\n"
+        else:
+            yield f'data: {{"chunk": "No file uploaded for analysis."}}\n\n'
+            yield "data: [DONE]\n\n"
+    return gen
 
 def check_issues(files):
-    if files:
-        text = extract_text_from_pdf(files[0])
-        if text:
-            issues = ask_gpt4o(f"Check for red flags, unusual clauses, or potential issues in this legal document: {text[:10000]}")
-            return f"Issues: {issues}"
-        return "No text extracted from PDF."
-    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 = """
-# ... (your CSS remains the same, omitted for brevity)
-"""
-
-theme = gr.themes.Base(
-    primary_hue="gray",
-    secondary_hue="gray",
-    neutral_hue="gray",
-).set(
-    body_text_color="#000000",
-    background_fill_primary="#6D8299"  # Slate blue background
-)
-
-with gr.Blocks(css=css, theme=theme, title="VerdictAI - Legal Assistant") as app:
-    jurisdiction = gr.State("KY")
-    chatbot = gr.Chatbot(elem_id="chat-container", label="Chat")
-    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"
-    )
-    with gr.Row():
-        with gr.Column(scale=2, elem_classes=["main-content"]):
-            chatbot
-            with gr.Row(elem_id="chat-input"):
-                msg
-                file_upload = gr.File(file_count="multiple", file_types=[".pdf", "image", "text"], 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], outputs=[chatbot, chatbot])
-    google_search_btn.click(fn=chat_interface, inputs=[msg, file_upload, chatbot, gr.State(True), jurisdiction], 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)
+    def gen():
+        if files:
+            texts = [extract_text_from_file(f) for f in files]
+            text = "\n".join(texts)
+            if text:
+                issues = ask_grok([{"role": "user", "content": f"Check for red flags, unusual clauses, or potential issues in this legal document(s) and highlight them: {text[:10000]}"}], stream=False)  # Explicitly non-stream
+                full_response = f"Highlighted Issues: {issues}"
+                chunks = [full_response[i:i+200] for i in range(0, len(full_response), 200)]
+                for part in chunks:
+                    yield f'data: {{"chunk": {json.dumps(part)}}}\n\n'
+                yield "data: [DONE]\n\n"
+            else:
+                yield f'data: {{"chunk": "No text extracted from file."}}\n\n'
+                yield "data: [DONE]\n\n"
+        else:
+            yield f'data: {{"chunk": "No file uploaded to check."}}\n\n'
+            yield "data: [DONE]\n\n"
+    return gen
+
+# Error handlers to always return JSON
+@app.errorhandler(400)
+def bad_request(error):
+    return jsonify({'error': 'Bad request'}), 400
+
+@app.errorhandler(404)
+def not_found(error):
+    return jsonify({'error': 'Not found'}), 404
+
+@app.errorhandler(405)
+def method_not_allowed(error):
+    return jsonify({'error': 'Method not allowed'}), 405
+
+@app.errorhandler(500)
+def internal_error(error):
+    return jsonify({'error': 'Internal server error'}), 500
+
+@app.errorhandler(Exception)
+def handle_exception(e):
+    logger.error(f"Unhandled exception: {str(e)}")
+    return jsonify({'error': str(e)}), 500
+
+# Flask routes
+@app.route('/')
+def index():
+    return send_from_directory('.', 'index.html')
+
+@app.route('/api/chat', methods=['POST'])
+def api_chat():
+    temp_paths = []  # Initialize here for finally block
+    def generate():
+        try:
+            # Early check for missing data
+            if 'payload' not in request.form:
+                yield f'data: {{"error": "Missing payload in request"}}\n\n'
+                yield "data: [DONE]\n\n"
+                return
+            payload = json.loads(request.form['payload'])
+            messages = payload['messages']
+            jurisdiction = payload['jurisdiction']
+            irac_mode = payload['irac_mode']
+            search_web = payload['web_search']
+            uploaded_files = request.files.getlist('file')
+            file_texts = []
+            if uploaded_files:
+                for file in uploaded_files:
+                    if file.filename:
+                        filename = secure_filename(file.filename)
+                        temp_path = os.path.join('/tmp', filename)
+                        file.save(temp_path)
+                        file_text = extract_text_from_file(temp_path)
+                        file_texts.append(file_text)
+                        temp_paths.append(temp_path)
+            file_text_combined = "\n".join(file_texts)
+            prompt = messages[-1]['content']  # for classification
+            task_type = classify_prompt(prompt)
+            if irac_mode:
+                task_type = "irac"
+            # Append file text to last user message if present
+            if file_text_combined:
+                messages[-1]['content'] += "\nAttached file content(s): " + file_text_combined[:10000]
+            if "summarize" in prompt.lower():
+                task_type = "document_analysis"
+                gen_func = summarize_document(temp_paths)
+                for chunk in gen_func():
+                    yield chunk
+            elif "analyze" in prompt.lower():
+                task_type = "document_analysis"
+                gen_func = analyze_document(temp_paths)
+                for chunk in gen_func():
+                    yield chunk
+            elif "check" in prompt.lower() or "issues" in prompt.lower() or "highlight" in prompt.lower():
+                task_type = "document_analysis"
+                gen_func = check_issues(temp_paths)
+                for chunk in gen_func():
+                    yield chunk
+            elif "generate" in prompt.lower() or "draft" in prompt.lower():
+                task_type = "document_creation"
+                for line in route_model(messages, task_type, temp_paths, search_web, jurisdiction):
+                    yield line
+            else:
+                for line in route_model(messages, task_type, temp_paths, search_web, jurisdiction):
+                    yield line
+            logger.info("Grok response streamed.")
+        except Exception as e:
+            logger.error(f"Error in /api/chat: {str(e)}")
+            yield f'data: {{"error": "{str(e)}"}}\n\n'
+            yield "data: [DONE]\n\n"
+        finally:
+            # Cleanup temp files (no context needed for os.remove)
+            for temp_path in temp_paths:
+                try:
+                    os.remove(temp_path)
+                except Exception as cleanup_e:
+                    logger.error(f"Cleanup error: {str(cleanup_e)}")
+    return Response(stream_with_context(generate()), mimetype='text/event-stream')
+
+@app.route('/download/<filename>', methods=['GET'])
+def download(filename):
+    return send_file(os.path.join('/tmp', filename), as_attachment=True)
+
+@app.route('/health', methods=['GET'])
+def health():
+    return "OK", 200
+
+if __name__ == '__main__':
+    logger.info("✅ All init complete. Starting Flask app...")
+    print("Flask app starting...")  # Fallback print
+    app.run(host='0.0.0.0', port=7860)