Update app.py

app.py

@@ -1,13 +1,15 @@
 # app.py
 # This is the updated main script. Copy-paste this over your existing app.py.
 # Changes:
-# - Removed unused 'bm25s' import (replaced with rank_bm25 in retrieval.py).
-# - Integrated retrieve_context from retrieval.py, which now uses hybrid_cap_search with lazy loading.
-# - Added handling for missing CAP components, logging warnings and skipping RAG if caches are absent.
-# - Retained Flask for serving the custom HTML+CSS+JS frontend and API endpoint /api/chat.
-# - Kept file handling, IRAC mode, web search toggle, and task classification logic.
-# - Updated route_model to use retrieve_context only if CAP components are available.
-# - Note: Precompute CAP components with precompute_cap_embeddings.py before deployment.
+# - Fixed import for OpenAI (added line break and ensured it's not commented out).
+# - Retained 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).
+# - Forced task_type to "irac" if IRAC mode is enabled; otherwise, uses classify_prompt.
+# - Forced web_search toggle handling.
+# - Forced task to "document_creation" routes directly to the fine-tuned GPT model.
+# - Updated route_model to use retrieve_context(prompt, task_type) instead of separate semantic_search/municipal_search.
+# - For document_creation/summaries, skip RAG (no retrieve_context call) to avoid slowdown.
 
 from flask import Flask, request, jsonify, send_from_directory
 from werkzeug.utils import secure_filename
@@ -18,11 +20,12 @@ import pdfplumber
 from docx import Document
 from googleapiclient.discovery import build
 import re
-from retrieval import retrieve_context, municipal_search  # Updated import
+from retrieval import retrieve_context  # Import from retrieval.py
 from task_processing import process_task_response
 from gpt_helpers import ask_gpt41_mini
 from prompt_builder import build_saul_prompt, build_editor_prompt
 from post_processing import ground_statutes
+from openai import OpenAI  # Fixed import for OpenAI client
 
 app_flask = Flask(__name__)
 os.environ["HF_HOME"] = "/data/.huggingface"
@@ -76,22 +79,18 @@ STATES = {
 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"] and not os.getenv("SKIP_CAP_INIT", "false").lower() == "true":
-        # Check if CAP components are available
-        if all(os.path.exists(f"/data/cap_{ext}") for ext in ["tfidf.pkl", "tfidf_matrix.npz", "gte.npy", "openai.npy"]):
-            combined_results = retrieve_context(prompt, task_type)
-            # Filter by jurisdiction if specified
-            if jurisdiction and jurisdiction != "All States":
-                state_name = STATES.get(jurisdiction, "")
-                state_code = jurisdiction
-                combined_results = [r for r in combined_results if any(s in (r.get('citation', '') + r.get('name', '')) for s in [state_code, state_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}"
-        else:
-            logger.warning("CAP hybrid components missing. Precompute them with precompute_cap_embeddings.py. Skipping RAG.")
-    
+    if task_type in ["case_law", "irac", "statute"]:
+        combined_results = retrieve_context(prompt, task_type)
+        # Filter by jurisdiction if specified
+        if jurisdiction and jurisdiction != "All States":
+            state_name = STATES.get(jurisdiction, "")
+            state_code = jurisdiction
+            combined_results = [r for r in combined_results if any(s in (r.get('citation', '') + r.get('name', '')) for s in [state_code, state_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}"
+
     if task_type == "document_creation":
         # Route directly to fine-tuned GPT for document creation
         saul_response = ask_gpt41_mini(prompt, jurisdiction)
@@ -176,7 +175,7 @@ def extract_text_from_file(file_path):
 def classify_prompt(prompt):
     prompt_lower = prompt.lower()
     if "summarize" in prompt_lower:
-        return "document_analysis"
+        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"]):