Update task_processing.py

task_processing.py

@@ -1,48 +1,83 @@
-# task_processing.py (Fully Updated - Minor Changes for Artifact Compatibility)
-# Added optional artifact generation in process_task_response if task involves documents.
-# Minor changes: Removed retrieve_context (aligned with retrieval.py updates); set rag_context to empty.
-# Avoided redundant GPT calls; focused on post-processing Grok's response (e.g., formatting).
+# task_processing.py (Updated for new VerdictAI modes)
+# - Adds labels/headers for new modes introduced in app.py
+# - Keeps lightweight, post-processing-only behavior
+# - No redundant model calls (no extra Grok/GPT invocations)
+# - Safe default passthrough if a mode/task_type is unrecognized
 
-from gpt_helpers import ask_gpt41_mini  # Retained for document_creation draft (but note redundancy with app.py)
+# Retained import in case you reference it elsewhere; not used here.
+from gpt_helpers import ask_gpt41_mini  # noqa: F401
 
-def process_task_response(task_type, grok_response, prompt, jurisdiction):
+
+def _wrap(title: str, jurisdiction: str, content: str) -> str:
+    """Minimal, consistent wrapper for returning labeled content."""
+    jur = f" ({jurisdiction})" if jurisdiction else ""
+    return f"{title}{jur}:\n\n{content}"
+
+
+def process_task_response(task_type: str, grok_response: str, prompt: str, jurisdiction: str) -> str:
     """
-    Process the Grok response based on task type. Since rag_context is already prepended in app.py,
-    this focuses on minimal post-processing (e.g., formatting headers) without redundant retrieval.
-    If additional context is needed (rare), use retrieve_context.
+    Post-process Grok output for display/download. This function is intentionally lightweight:
+    - No additional retrieval or model calls
+    - Just adds consistent headings or minimal formatting per task type
+    - Safe passthrough for unknown types
     """
-    # Fetch context only if absolutely needed for post-processing (e.g., if it's already in the prompt
-    rag_context = ""  # Default to empty; override per task if needed (no more retrieve_context)
-    if task_type == "case_law" or task_type == "irac":
-        # No need to prepend; Grok already has rag_context with full case text/citations
-        return f"Case Law / IRAC Analysis ({jurisdiction}): {grok_response}"
-    elif task_type == "statute":
-        # No need to prepend; Grok has municipal data from retrieve_context
-        return f"Statute Analysis ({jurisdiction}): {grok_response}"
-    elif task_type == "document_creation":
-        # app.py already handles draft + polish; if post-processing needed, format only
-        # Avoid calling ask_gpt41_mini again to prevent duplication
-        return f"Final Document ({jurisdiction}): {grok_response}"
-    elif task_type == "document_analysis":
-        return f"Document Analysis ({jurisdiction}): {grok_response}"
-    elif task_type == "analogical_reasoning":
-        return f"Analogical Reasoning ({jurisdiction}): {grok_response}"
-    elif task_type == "deposition_questions":
-        return f"Deposition Questions ({jurisdiction}): {grok_response}"
-    elif task_type == "verdict_predictor":
-        return f"Verdict Prediction ({jurisdiction}): {grok_response}"
-    elif task_type == "damages_estimator":
-        return f"Damages Estimate ({jurisdiction}): {grok_response}"
-    elif task_type == "pleading_validator":
-        return f"Pleading Validation ({jurisdiction}): {grok_response}"
-    elif task_type == "legal_inbox":
-        return f"Legal Inbox Processing ({jurisdiction}): {grok_response}"
-    elif task_type == "privilege_analyzer":
-        return f"Privilege/PII Analysis ({jurisdiction}): {grok_response}"
-    elif task_type == "eli5":
-        return f"ELI5 Explanation ({jurisdiction}): {grok_response}"
-    elif task_type == "workflow_template":
-        return f"Workflow Template ({jurisdiction}): {grok_response}"
-    elif task_type == "legal_strategy":
-        return f"Legal Strategy ({jurisdiction}): {grok_response}"
-    return grok_response
+    tt = (task_type or "").lower()
+
+    # Legacy/primary types
+    if tt in ("case_law", "irac"):
+        return _wrap("Case Law / IRAC Analysis", jurisdiction, grok_response)
+    if tt == "statute":
+        return _wrap("Statute Analysis", jurisdiction, grok_response)
+    if tt == "document_creation":
+        return _wrap("Final Document", jurisdiction, grok_response)
+    if tt == "document_analysis":
+        return _wrap("Document Analysis", jurisdiction, grok_response)
+    if tt == "legal_strategy":
+        return _wrap("Legal Strategy", jurisdiction, grok_response)
+
+    # Existing misc types (kept for compatibility with any callers you might have)
+    if tt == "analogical_reasoning":
+        return _wrap("Analogical Reasoning", jurisdiction, grok_response)
+    if tt in ("deposition_questions", "deposition_qs"):
+        return _wrap("Deposition Questions", jurisdiction, grok_response)
+    if tt == "verdict_predictor":
+        return _wrap("Verdict Prediction", jurisdiction, grok_response)
+    if tt == "damages_estimator":
+        return _wrap("Damages Estimate", jurisdiction, grok_response)
+    if tt == "pleading_validator":
+        return _wrap("Pleading Validation", jurisdiction, grok_response)
+    if tt == "legal_inbox":
+        return _wrap("Legal Inbox Processing", jurisdiction, grok_response)
+    if tt == "privilege_analyzer":
+        return _wrap("Privilege/PII Analysis", jurisdiction, grok_response)
+    if tt == "eli5":
+        return _wrap("Plain-English Explanation", jurisdiction, grok_response)
+    if tt == "workflow_template":
+        return _wrap("Workflow Template", jurisdiction, grok_response)
+
+    # New modes introduced in app.py
+    if tt == "citation_check":
+        return _wrap("Citation Check / Shepardizing-lite", jurisdiction, grok_response)
+    if tt == "brief_builder":
+        return _wrap("Brief Builder (IRAC Sections)", jurisdiction, grok_response)
+    if tt in ("deposition_qa",):
+        return _wrap("Deposition Q&A Generator", jurisdiction, grok_response)
+    if tt == "checklist":
+        return _wrap("Checklist & Red Flags", jurisdiction, grok_response)
+    if tt == "comparative_law":
+        return _wrap("Comparative Law", jurisdiction, grok_response)
+    if tt == "client_summary":
+        return _wrap("Client-Friendly Summary", jurisdiction, grok_response)
+    if tt == "motion_skeleton":
+        return _wrap("Motion Skeleton", jurisdiction, grok_response)
+    if tt == "opposing_argument":
+        return _wrap("Opposing Argument (Devil's Advocate)", jurisdiction, grok_response)
+    if tt == "case_extractor":
+        return _wrap("Case/Citation Extractor", jurisdiction, grok_response)
+    if tt == "judge_style":
+        return _wrap("Judge/County Style Adaptation", jurisdiction, grok_response)
+    if tt == "general_qa":
+        return _wrap("General Q&A", jurisdiction, grok_response)
+
+    # Fallback passthrough
+    return grok_response