Update app.py

app.py

@@ -10,7 +10,6 @@ import openai
 import tempfile
 import json
 import re
-
 import markdown
 from bs4 import BeautifulSoup
 
@@ -121,55 +120,81 @@ def markdown_to_docx(md_text: str) -> str:
         elif el.name == "li":
             doc.add_paragraph(f"• {el.get_text()}")
 
-    # Save into a temporary file
     tmp_path = os.path.join(tempfile.gettempdir(), "draft.docx")
     doc.save(tmp_path)
     return tmp_path
-    
-# ----------------- MAIN FUNCTION -----------------
+
+# ----------------- DRAFT ANALYZER -----------------
+def analyze_template_draft(ref_text: str) -> str:
+    """Analyze uploaded draft structure, tone, style, headings, and expected content."""
+    if not ref_text:
+        return "(no template provided)"
+
+    system_prompt = """You are a legal draft analyzer. 
+Your task is to carefully analyze the uploaded legal draft document and summarize its full structure and style. 
+Extract the following information clearly and systematically:
+1. Headings and subheadings (exact order).
+2. Approximate length/word count per section.
+3. Purpose of each section (what content it usually contains).
+4. Writing style and tone (formal/informal, persuasive, assertive, etc.).
+5. Formatting conventions (headings, numbering, bullet points, capitalization).
+6. Sentence/paragraph length and complexity.
+7. Any special legal phrases or terminology patterns.
+8. Any notes on length and overall flow.
+
+Return a structured MARKDOWN report that can be given as instructions to another model.
+Do not rewrite the draft, only analyze it."""
+
+    try:
+        resp = openai_client.chat.completions.create(
+            model="gpt-4o-mini",
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": ref_text[:40000]}  # limit for context
+            ],
+            temperature=0.2,
+            max_tokens=4000
+        )
+        analysis = resp.choices[0].message.content.strip()
+    except Exception as e:
+        analysis = f"(Analyzer error: {e})"
+
+    return analysis
+
+# ----------------- MAIN -----------------
 def generate_legal_draft(case_text, uploaded_file=None, add_citations=True):
-    ref_text = load_reference_text(uploaded_file)
+    # Step 0: Show initial status
+    yield gr.update(value="🔍 Searching in Knowledge Base..."), None
 
     queries = build_queries_with_llm(case_text)
     context_text, citations = pinecone_search(queries, top_k=10)
 
-    system_prompt = """You are an expert legal drafter for Pakistani law. Your task is to create a professional, court-ready legal petition in MARKDOWN format using three inputs:
+    yield gr.update(value="📝 Analyzing Template Document..."), None
+    ref_text = load_reference_text(uploaded_file)
+    template_analysis = analyze_template_draft(ref_text)
+
+    yield gr.update(value="⚖️ Generating Final Draft..."), None
+    system_prompt = """You are an expert legal drafter for Pakistani law. Your task is to create a professional, court-ready legal petition in MARKDOWN format using four inputs:
                         1. User Input: Case details including client info, petition type, court, facts, relevant laws, and sections.
                         2. Knowledge Base Context: Relevant laws, case precedents, and ordinances retrieved from the vector database (Constitution of Pakistan, Punjab case law, FBR ordinances).
-                        3. Reference Template: A legal document uploaded by the user. Extract its **structure, headings, section order, style, tone, and formatting** — but do not copy its text.
+                        3. Template Draft Analysis: A structured analysis of an uploaded legal template (headings, section purposes, tone, formatting rules, length, style).
+                        4. Fallback: If some info is missing, state explicitly instead of hallucinating.
                         Instructions
-                        1. Structure & Headings
-                        - Replicate the section hierarchy of the uploaded template.
-                        - Main heading should represent the petition title or case name.
-                        - Subheadings and any lower-level headings must match the style and order from the template.
-                        - Sections should include: Parties, Facts, Legal Grounds, Arguments, Prayer, etc., as per the template.
-                        2. Tone & Style
-                        - Use formal, professional, and persuasive legal language.
-                        - Facts are objective; legal arguments are assertive.
-                        - Follow the tone and phrasing style of the uploaded template document.
-                        3. Content Integration
-                        - Incorporate relevant context from the vector database where appropriate.
-                        - Cite legal provisions clearly in-text when relevant.
-                        - Ensure content is logically coherent, comprehensive, and supports the petition’s objective.
-                        - Do not hallucinate laws or precedents.
-                        4. References
-                        - Include a "References" section at the end if citations are present.
-                        - Format as: `1. Source Name (score)`
-                        5. Output Rules
-                        - Produce output MARKDOWN.
-                        - Do not add explanations, summaries, or extra text.
-                        - Maintain clarity, professionalism, and adherence to legal drafting standards.
-                        - Preserve structure, tone, style, and headings from the uploaded template as much as possible.
-                        6. Fallback
-                        - If context or relevant laws are missing, state: "No applicable precedent found" or "Relevant law cited above."
+                        - Replicate the section hierarchy and style described in the template analysis.
+                        - Ensure clarity, professionalism, and persuasive legal argumentation.
+                        - Integrate legal context where appropriate with accurate citations.
+                        - Output must be MARKDOWN only, no explanations or extra commentary.
 """
+
     user_prompt = f"""
 **User Input:** 
 {case_text}
+
 **Knowledge Base Context:**
 {context_text or '(no matches)'}
-**Reference Template (format+headings+structure+tone+language):**
-{(ref_text[:50000] + '...') if ref_text else '(none)'}
+
+**Template Draft Analysis:**
+{template_analysis}
 """
 
     try:
@@ -180,7 +205,7 @@ def generate_legal_draft(case_text, uploaded_file=None, add_citations=True):
                 {"role":"user","content":user_prompt}
             ],
             max_completion_tokens=15000,
-            verbosity= "high"
+            verbosity="high"
         )
         draft_md = resp.choices[0].message.content.strip()
     except Exception as e:
@@ -192,7 +217,7 @@ def generate_legal_draft(case_text, uploaded_file=None, add_citations=True):
             draft_md += f"{i}. {c['source']} (score: {c['score']:.3f})\n"
 
     docx_path = markdown_to_docx(draft_md)
-    return gr.update(value=draft_md), markdown_to_docx(draft_md)
+    yield gr.update(value=draft_md), markdown_to_docx(draft_md)
 
 
 
@@ -207,16 +232,13 @@ with gr.Blocks() as demo:
     draft_output = gr.Markdown(label="Draft Output")
 
     download_btn = gr.DownloadButton(label="⬇️ Download Word")
-    # draft_output1 = gr.Markdown("# HI")
 
     btn = gr.Button("Generate Draft")
     btn.click(
-        generate_legal_draft,
-        inputs=[case_text, uploaded_file, add_citations],
-        outputs=[draft_output, download_btn],
+    generate_legal_draft,
+    inputs=[case_text, uploaded_file, add_citations],
+    outputs=[draft_output, download_btn]
 )
-
-
 # ----------------- RUN -----------------
-if __name__ == "__main__":
-    demo.launch()
+if __name__ == "__main__": 
+    demo.launch()