Update app.py

app.py

@@ -19,18 +19,14 @@ from requests.adapters import HTTPAdapter
 from http.client import HTTPConnection
 from gpt_helpers import ask_gpt41_mini
 from prompt_builder import build_grok_prompt, build_editor_prompt
-
 # Enable HTTP debugging for detailed request/response logging
 HTTPConnection.debuglevel = 1
-
 app = Flask(__name__)
 os.environ["HF_HOME"] = "/data/.huggingface"
-
 # Logging setup
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger("app")
 logger.debug("✅ Logging initialized. Starting app setup.")
-
 # File storage setup
 UPLOAD_DIR = '/data/uploads'
 try:
@@ -39,20 +35,16 @@ try:
     logger.debug(f"Created/verified upload directory: {UPLOAD_DIR}")
 except Exception as e:
     logger.error(f"Failed to create upload directory {UPLOAD_DIR}: {str(e)}")
-
 uploaded_files = []
 file_lock = threading.Lock()
-
 # Grok API setup (intentionally hard-coded per user's request)
 GROK_API_URL = "https://api.x.ai/v1/chat/completions"
 GROK_API_TOKEN = "xai-xHzXVZ2TtRwcqP0UZO8e0fCTtuvg9ZKfHYkdBflYVLxcXHLUbUXQbULs3N5TxcxqzE8E7Yy0cK2umCON"
-
 # OpenAI key (optional path you already had)
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAI_API_KEY_VERDICTAI")
 if not OPENAI_API_KEY:
     logger.error("OPENAI_API_KEY or OPENAI_API_KEY_VERDICTAI not set in environment variables.")
 logger.debug("✅ Grok and OpenAI API endpoints and tokens set.")
-
 # Define STATES dictionary
 STATES = {
     "KY": "Kentucky",
@@ -63,15 +55,12 @@ STATES = {
     "TN": "Tennessee",
     # Add more as needed
 }
-
 # Global citation regex (more permissive to catch common formats like S.W.3d, F.3d, etc.)
 CITATION_RE = re.compile(
     r'([A-Z][A-Za-z&.\- ]+ v\. [A-Z][A-Za-z&.\- ]+, \d+ [A-Z][A-Za-z.]*\d* \d+ \([A-Za-z. ]+ \d{4}\))'
 )
-
 def extract_citations(text: str):
     return CITATION_RE.findall(text or "")
-
 def extract_text_from_file(file_path):
     try:
         if not os.path.exists(file_path):
@@ -95,7 +84,6 @@ def extract_text_from_file(file_path):
     except Exception as e:
         logger.error(f"Error extracting text from {file_path}: {str(e)}")
         return f"Error: Failed to extract text from {file_path} - {str(e)}"
-
 def classify_prompt(prompt):
     """Baseline classifier when no explicit 'mode' is provided."""
     p = prompt.lower()
@@ -132,7 +120,6 @@ def classify_prompt(prompt):
     if any(x in p for x in ["summarize", "what does", "order"]):
         return "document_analysis"
     return "general_qa"
-
 def ask_grok(messages, stream=False, deep_search=False):
     try:
         headers = {
@@ -222,7 +209,6 @@ def ask_grok(messages, stream=False, deep_search=False):
                 yield "data: [DONE]\n\n"
             return error_gen()
         return "[Grok Error] " + str(e)
-
 def _set_doc_defaults(doc):
     style = doc.styles['Normal']
     style.font.name = 'Times New Roman'
@@ -233,7 +219,6 @@ def _set_doc_defaults(doc):
         sec.bottom_margin = Inches(1)
         sec.left_margin = Inches(1.25)
         sec.right_margin = Inches(1)
-
 def _add_caption_block(doc, court_line, action_no, plaintiff, defendant, motion_title):
     p = doc.add_paragraph()
     p.alignment = WD_ALIGN_PARAGRAPH.CENTER
@@ -256,7 +241,6 @@ def _add_caption_block(doc, court_line, action_no, plaintiff, defendant, motion_
     ]:
         run = p.add_run(line + "\n")
         run.bold = True if line.isupper() and line not in ("PLAINTIFF,", "DEFENDANT,", "v.") else False
-
 def create_legal_docx(content, jurisdiction, filename, task_type):
     try:
         doc = Document()
@@ -302,7 +286,6 @@ def create_legal_docx(content, jurisdiction, filename, task_type):
     except Exception as e:
         logger.error(f"Error creating DOCX {filename}: {str(e)}")
         return f"Error creating document: {str(e)}"
-
 def web_search(q: str) -> str:
     try:
         url = f"https://www.google.com/search?q={requests.utils.quote(q)}"
@@ -327,15 +310,12 @@ def web_search(q: str) -> str:
     except Exception as e:
         logger.error(f"Web search error: {str(e)}")
         return ""
-
 # ---------- Feature Prompt Builders ----------
-
 def build_prompt_for_mode(mode, user_prompt, jurisdiction, state_name, combined_text):
     """Return (system_message, user_message) for a given mode."""
     base_context = f"Jurisdiction: {state_name} ({jurisdiction}). Assume Kentucky practice unless otherwise stated. Write in clear, professional legal prose tailored to Kentucky courts."
     text_note = f"\n\nAttached/Extracted Documents (truncated):\n{combined_text[:10000]}" if combined_text else ""
     mode = (mode or "").lower()
-
     if mode == "citation_check":
         # Use any explicit citations in user prompt + docs
         all_text = (user_prompt or "") + "\n" + (combined_text or "")
@@ -344,70 +324,57 @@ def build_prompt_for_mode(mode, user_prompt, jurisdiction, state_name, combined_
         sys = base_context + " Verify legal citations and report treatment (followed, distinguished, limited, overruled) with brief parentheticals."
         usr = f"Please verify the following citations and list significant subsequent cases and treatments:\n{cit_block}\n{text_note}"
         return sys, usr
-
     if mode == "brief_builder":
         sys = base_context + " Build a structured brief with clear headings."
         usr = f"Using the facts and issues below, produce a Kentucky-style brief with sections: FACTS, ISSUE(S), RULE(S), APPLICATION/ANALYSIS, CONCLUSION.\n\nUser Prompt:\n{user_prompt}\n{text_note}"
         return sys, usr
-
     if mode == "deposition_qa":
         sys = base_context + " Generate practical, targeted deposition questions."
         usr = f"Based on the matter below, draft 20 deposition questions for (1) the opposing party, (2) a key fact witness, and (3) an expert (if applicable). Group by witness type and topic.\n\n{user_prompt}\n{text_note}"
         return sys, usr
-
     if mode == "checklist":
         sys = base_context + " Create a compliance checklist and red-flag analysis."
         usr = f"Produce a Kentucky practice checklist for this document or scenario, and list red flags or missing essentials (caption, certificate of service, venue, signatures, deadlines, rules).\n\n{user_prompt}\n{text_note}"
         return sys, usr
-
     if mode == "comparative_law":
         sys = base_context + " Compare Kentucky with Indiana, Ohio, and Federal law succinctly."
         usr = f"Answer under Kentucky law first. Then briefly compare with Indiana, Ohio, and Federal law, noting major similarities/differences and key citations.\n\nQuestion:\n{user_prompt}\n{text_note}"
         return sys, usr
-
     if mode == "client_summary":
         sys = base_context + " Translate into plain, client-friendly English. Avoid legalese."
         usr = f"Summarize in plain English suitable for a client. Provide bullets of key points, obligations, and risks.\n\n{user_prompt}\n{text_note}"
         return sys, usr
-
     if mode == "motion_skeleton":
         sys = base_context + " Draft a Kentucky-style motion skeleton with placeholders."
         usr = f"Draft a motion skeleton for '{user_prompt}'. Include: Caption placeholder, Introduction, Legal Standard (with rule/case placeholders), Argument headings, Relief Requested, Signature block, Certificate of Service. Use placeholders like [Insert Facts].\n{text_note}"
         return sys, usr
-
     if mode == "opposing_argument":
         sys = base_context + " Write the strongest opposing argument under Kentucky law."
         usr = f"Given the following draft or position, write the best counter-argument opposing counsel would make. Provide concise headings and authorities.\n\n{user_prompt}\n{text_note}"
         return sys, usr
-
     if mode == "case_extractor":
         sys = base_context + " Extract and list all legal citations from the provided text."
         usr = f"Extract every legal citation found and list them one per line. Then group by jurisdiction (KY, Federal, Other). If context suggests relevance, add a 1-2 line parenthetical per citation.\n\n{user_prompt}\n{text_note}"
         return sys, usr
-
     if mode == "judge_style":
         sys = base_context + " Adapt tone and format to the specified judge/county preferences."
         usr = f"Draft the requested text in the style preferred by the specified judge or county as described below.\n\n{user_prompt}\n{text_note}"
         return sys, usr
-
     # Fallback: use existing prompt builder pathways
     sys = build_grok_prompt(user_prompt, mode or "general_qa", jurisdiction, "")
     usr = f"{user_prompt}\n{text_note}"
     return sys, usr
-
 # ---------- Pipeline Stages ----------
-
 def route_model(messages, task_type, files, deep_search, jurisdiction, explicit_mode=None):
     logger.debug(f"Starting route_model with task_type: {task_type}, jurisdiction: {jurisdiction}, deep_search: {deep_search}, files: {files}, explicit_mode: {explicit_mode}")
     try:
         jurisdiction = jurisdiction if jurisdiction in STATES else "KY"
         state_name = STATES.get(jurisdiction, "Kentucky")
-        rag_context = ""  # reserved for future RAG
+        rag_context = "" # reserved for future RAG
         needs_deep_search = (
             (task_type not in ["document_creation"])
             and (deep_search or "research" in messages[-1]['content'].lower())
         )
-
         yield f'data: {{"chunk": "Step 1: Analyzing query..."}}\n\n'
         prompt = messages[-1]['content']
         prompt_lower = prompt.lower()
@@ -415,16 +382,13 @@ def route_model(messages, task_type, files, deep_search, jurisdiction, explicit_
             if len(messages) > 1 and messages[-2]['role'] == 'assistant':
                 prior_summary = messages[-2]['content'][:500] + "..."
                 messages[-1]['content'] += f"\n\nPrior context to retry: {prior_summary}"
-
         # Gather file text
         file_text_combined = "\n".join([extract_text_from_file(p) for p in files if p])
         if file_text_combined and not file_text_combined.startswith("Error"):
             logger.debug(f"Added document content to prompt: {file_text_combined[:200]}...")
         messages[-1]['content'] = f"{messages[-1]['content']}\n\n{rag_context}"
-
         # Branching by task type / explicit mode
         mode = explicit_mode or task_type
-
         # Document creation (kept as in your pipeline: GPT-4.1-mini draft -> Grok polish -> DOCX/PDF)
         if task_type in ["document_creation", "irac", "case_law", "statute", "legal_strategy", "general_qa"] and mode == task_type:
             yield f'data: {{"chunk": "Step 2: Generating initial draft with GPT-4.1-mini..."}}\n\n'
@@ -456,7 +420,6 @@ def route_model(messages, task_type, files, deep_search, jurisdiction, explicit_
                 logger.error(error_msg)
                 full_response = gpt_response
                 yield f'data: {{"chunk": "Grok failed, using GPT-4.1-mini draft..."}}\n\n'
-
         elif mode in [
             "citation_check", "brief_builder", "deposition_qa", "checklist",
             "comparative_law", "client_summary", "motion_skeleton",
@@ -472,7 +435,6 @@ def route_model(messages, task_type, files, deep_search, jurisdiction, explicit_
                 yield f'data: {{"error": {json.dumps(error_msg)}}}\n\n'
                 yield "data: [DONE]\n\n"
                 return
-
         else:
             # Fallback: single Grok pass with build_grok_prompt
             yield f'data: {{"chunk": "Step 2: Processing with Grok..."}}\n\n'
@@ -485,7 +447,6 @@ def route_model(messages, task_type, files, deep_search, jurisdiction, explicit_
                 yield f'data: {{"error": {json.dumps(error_msg)}}}\n\n'
                 yield "data: [DONE]\n\n"
                 return
-
         # Step 4: Verify citations lightly (text modes only)
         yield f'data: {{"chunk": "Step 4: Verifying citations and facts..."}}\n\n'
         citations = extract_citations(full_response)
@@ -510,17 +471,15 @@ def route_model(messages, task_type, files, deep_search, jurisdiction, explicit_
                 system_content = build_grok_prompt(prompt, task_type, jurisdiction, rag_context)
                 re_messages = [{'role': 'system', 'content': system_content}, {'role': 'user', 'content': re_prompt}]
             full_response = ask_grok(re_messages, stream=False, deep_search=True)
-            if full_response.startswith("[Grok Error"]:
+            if full_response.startswith("[Grok Error"):
                 error_msg = f"Grok re-polish failed: {full_response}"
                 logger.error(error_msg)
                 yield f'data: {{"error": {json.dumps(error_msg)}}}\n\n'
                 yield "data: [DONE]\n\n"
                 return
-
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
         pdf_filename = f"/tmp/legal_doc_{timestamp}.pdf"
         docx_filename = f"/tmp/legal_doc_{timestamp}.docx"
-
         if task_type == "document_creation" and mode == "document_creation":
             yield f'data: {{"chunk": "Step 5: Generating document and PDF preview..."}}\n\n'
             main_content = create_legal_docx(full_response, jurisdiction, docx_filename, task_type)
@@ -537,10 +496,9 @@ def route_model(messages, task_type, files, deep_search, jurisdiction, explicit_
                 error_msg = f"PDF conversion failed: {str(e)}"
                 logger.error(error_msg)
                 yield f'data: {{"error": {json.dumps(error_msg)}}}\n\n'
-                main_content = full_response  # still stream text
+                main_content = full_response # still stream text
         else:
             main_content = full_response
-
         yield f'data: {{"chunk": "Step 6: Finalizing response..."}}\n\n'
         chunks = [main_content[i:i+100] for i in range(0, len(main_content), 100)]
         for part in chunks:
@@ -553,9 +511,7 @@ def route_model(messages, task_type, files, deep_search, jurisdiction, explicit_
         logger.error(error_msg)
         yield f'data: {{"error": {json.dumps(error_msg)}}}\n\n'
         yield "data: [DONE]\n\n"
-
 # ---------- Utility Flows (unchanged behavior) ----------
-
 def summarize_document(files):
     def gen():
         try:
@@ -584,7 +540,6 @@ def summarize_document(files):
             yield f'data: {{"error": {json.dumps(error_msg)}}}\n\n'
             yield "data: [DONE]\n\n"
     return gen()
-
 def analyze_document(files):
     def gen():
         try:
@@ -613,7 +568,6 @@ def analyze_document(files):
             yield f'data: {{"error": {json.dumps(error_msg)}}}\n\n'
             yield "data: [DONE]\n\n"
     return gen()
-
 def check_issues(files):
     def gen():
         try:
@@ -642,14 +596,11 @@ def check_issues(files):
             yield f'data: {{"error": {json.dumps(error_msg)}}}\n\n'
             yield "data: [DONE]\n\n"
     return gen()
-
 # ---------- Flask Routes ----------
-
 @app.route('/')
 def index():
     logger.debug("Serving index.html")
     return send_from_directory('.', 'index.html')
-
 @app.route('/api/chat', methods=['POST'])
 def api_chat():
     logger.info("Received request to /api/chat")
@@ -664,30 +615,25 @@ def api_chat():
             message = data.get('message', '')
             if not message:
                 yield f'data: {{"error": "No message provided"}}\n\n'
-                yield "data: [DONE]\n\n"
+                yield "data: [DONE]\n\n'
                 return
-
             # Optional explicit mode from client (e.g., "citation_check", "deposition_qa", etc.)
-            explicit_mode = data.get('mode')  # if provided, we'll use it
+            explicit_mode = data.get('mode') # if provided, we'll use it
             jurisdiction = data.get('jurisdiction', 'KY')
             irac_mode = data.get('irac', False)
             deep_search = data.get('deepSearch', False)
             files_filenames = data.get('files', [])
-
             temp_paths = []
             with file_lock:
                 for f in uploaded_files:
                     if f['filename'] in files_filenames:
                         temp_paths.append(f['path'])
             yield f'data: {{"uploaded_files": {json.dumps(files_filenames)}}}\n\n'
-
             prompt_lower = message.lower()
             task_type = classify_prompt(message)
             if irac_mode:
                 task_type = "irac"
-
             logger.info(f"Task type: {task_type}, Explicit mode: {explicit_mode}, Jurisdiction: {jurisdiction}, Deep search: {deep_search}, Files: {files_filenames}")
-
             # Fast paths for utility flows if keywords + files present
             if "summarize" in prompt_lower and temp_paths:
                 for chunk in summarize_document(temp_paths):
@@ -703,7 +649,6 @@ def api_chat():
                 messages = [{'role': 'user', 'content': message}]
                 for line in route_model(messages, task_type, temp_paths, deep_search, jurisdiction, explicit_mode=explicit_mode):
                     yield line
-
             logger.info("Response streamed successfully.")
         except Exception as e:
             error_msg = f"Error in /api/chat: {str(e)}"
@@ -711,7 +656,6 @@ def api_chat():
             yield f'data: {{"error": {json.dumps(error_msg)}}}\n\n'
             yield "data: [DONE]\n\n"
     return Response(stream_with_context(generate()), mimetype='text/event-stream', headers={'Cache-Control': 'no-cache', 'Connection': 'keep-alive', 'X-Accel-Buffering': 'no'})
-
 @app.route('/api/files', methods=['POST'])
 def upload_files():
     logger.info("Received request to /api/files POST")
@@ -741,7 +685,6 @@ def upload_files():
     except Exception as e:
         logger.error(f"Error in upload_files: {str(e)}")
         return jsonify({'error': str(e)}), 500
-
 @app.route('/api/files', methods=['GET'])
 def list_files():
     logger.info("Received request to /api/files GET")
@@ -757,7 +700,6 @@ def list_files():
     except Exception as e:
         logger.error(f"Error in list_files: {str(e)}")
         return jsonify({'error': str(e)}), 500
-
 @app.route('/api/files/<filename>', methods=['DELETE'])
 def delete_file(filename):
     logger.info(f"Received request to delete file: {filename}")
@@ -777,7 +719,6 @@ def delete_file(filename):
     except Exception as e:
         logger.error(f"Error in delete_file: {str(e)}")
         return jsonify({'error': str(e)}), 500
-
 @app.route('/download/<filename>', methods=['GET'])
 def download(filename):
     logger.info(f"Received request to download: {filename}")
@@ -790,12 +731,10 @@ def download(filename):
     except Exception as e:
         logger.error(f"Error in download: {str(e)}")
         return jsonify({'error': str(e)}), 500
-
 @app.route('/health', methods=['GET'])
 def health():
     logger.debug("Health check requested")
     return jsonify({'status': 'healthy'}), 200
-
 if __name__ == '__main__':
     # Run the Flask app
-    app.run(host='0.0.0.0', port=7860, debug=True)
+    app.run(host='0.0.0.0', port=7860, debug=True)