App.py

# app.py
import requests
import os
import logging
from datetime import datetime
import pdfplumber
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, load_from_disk
from sentence_transformers import SentenceTransformer
import torch
import numpy as np
import shutil
from huggingface_hub import hf_hub_download
import pickle
import faiss
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

app = Flask(__name__)  # Renamed from app_flask to app for HF Spaces compatibility

os.environ["HF_HOME"] = "/data/.huggingface"

# 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

# Hugging Face authentication
from huggingface_hub import login
hf_token = os.environ.get("HF_TOKEN", "")
if hf_token:
    login(hf_token)
    logger.info("✅ Authenticated with Hugging Face token for gated repos.")
else:
    logger.warning("HF_TOKEN not set; gated repos may not be accessible.")

# Check environment variables
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:
    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

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"
}
logger.info("✅ States dictionary loaded.")

# Verdict Ai api call function (updated for streaming)
def ask_grok(messages, stream=False):
    try:
        headers = {
            "Accept": "application/json",
            "Content-Type": "application/json",
            "Authorization": f"Bearer {GROK_API_TOKEN}"
        }
        payload = {
            "messages": messages,
            "model": "grok-4-0709",
            "stream": stream,
            "temperature": 0.1,
            "max_tokens": 131072,  # High value for long responses
            "search_parameters": {
                "mode": "on"
            }
        }
        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()
        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:
            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"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:
        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"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", "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", "title opinion"]):
        return "document_creation"
    elif any(k in prompt_lower for k in ["review", "summarize", "clause", "red flags"]):
        return "document_analysis"
    elif any(k in prompt_lower for k in ["analogous", "similar", "pattern"]):
        return "analogical_reasoning"
    elif any(k in prompt_lower for k in ["deposition", "cross-exam", "questions"]):
        return "deposition_questions"
    elif any(k in prompt_lower for k in ["verdict", "outcome", "likely"]):
        return "verdict_predictor"
    elif any(k in prompt_lower for k in ["damages", "verdict amounts"]):
        return "damages_estimator"
    elif any(k in prompt_lower for k in ["validate", "motion", "complaint"]):
        return "pleading_validator"
    elif any(k in prompt_lower for k in ["sort", "timeline", "posture"]):
        return "legal_inbox"
    elif any(k in prompt_lower for k in ["privilege", "pii", "redact"]):
        return "privilege_analyzer"
    elif any(k in prompt_lower for k in ["eli5", "explain like", "plain english"]):
        return "eli5"
    elif any(k in prompt_lower for k in ["checklist", "plan", "template"]):
        return "workflow_template"
    elif any(k in prompt_lower for k in ["strategy", "pros", "cons"]):
        return "legal_strategy"
    return "general_qa"

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:
            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:
        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 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):
    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):
    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):
    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)