Update app.py

app.py

@@ -9,7 +9,6 @@ from pinecone import Pinecone
 import openai
 import json
 import re
-import markdown
 
 # ----------------- CONFIG -----------------
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
@@ -25,6 +24,8 @@ embedder = SentenceTransformer("all-mpnet-base-v2")
 
 # ----------------- HELPERS -----------------
 def load_reference_text(uploaded_file) -> str:
+    if not uploaded_file:
+        return ""
     if uploaded_file.name.lower().endswith(".docx"):
         doc = Document(uploaded_file)
         return "\n".join(p.text for p in doc.paragraphs)
@@ -38,8 +39,7 @@ def load_reference_text(uploaded_file) -> str:
         return text
     elif uploaded_file.name.lower().endswith(".txt"):
         return uploaded_file.read().decode("utf-8", errors="ignore")
-    else:
-        return ""
+    return ""
 
 def parse_json_safe(raw_text: str, fallback: str) -> List[str]:
     try:
@@ -54,7 +54,7 @@ def build_queries_with_llm(user_text: str, max_queries: int = 15) -> List[str]:
     system_prompt = (
         "You are a legal research assistant. "
         "A new petition needs to be drafted using the following client/case description. "
-        "Devise 5-6 or more concise queries that will be helpful to retrieve relevant information "
+        "Devise 5–6 or more concise queries that will be helpful to retrieve relevant information "
         "from a knowledge base containing the Constitution of Pakistan, Punjab case law, "
         "and FBR tax ordinances. "
         "Return ONLY a JSON array of strings, no extra text."
@@ -73,7 +73,7 @@ def build_queries_with_llm(user_text: str, max_queries: int = 15) -> List[str]:
         )
         raw = resp.choices[0].message.content.strip()
         queries = parse_json_safe(raw, user_text)
-    except Exception as e:
+    except Exception:
         queries = [user_text[:512]]
     return queries[:max_queries]
 
@@ -114,7 +114,7 @@ def markdown_to_docx(md_text: str) -> BytesIO:
 
 # ----------------- MAIN FUNCTION -----------------
 def generate_legal_draft(case_text, uploaded_file=None, add_citations=True):
-    ref_text = load_reference_text(uploaded_file) if uploaded_file else ""
+    ref_text = load_reference_text(uploaded_file)
 
     queries = build_queries_with_llm(case_text)
     context_text, citations = pinecone_search(queries, top_k=10)
@@ -157,6 +157,7 @@ Instructions
 **Reference Template (format+headings+structure+tone+language):**
 {(ref_text[:50000] + '...') if ref_text else '(none)'}
 """
+
     try:
         resp = openai_client.chat.completions.create(
             model="gpt-4o-mini",
@@ -176,23 +177,25 @@ Instructions
         for i, c in enumerate(citations, 1):
             draft_md += f"{i}. {c['source']} (score: {c['score']:.3f})\n"
 
-    return draft_md, (markdown_to_docx(draft_md), "legal_draft.docx")
+    return draft_md, markdown_to_docx(draft_md)
 
 # ----------------- GRADIO INTERFACE -----------------
-iface = gr.Interface(
-    fn=generate_legal_draft,
-    inputs=[
-        gr.Textbox(label="Case Details", lines=10, placeholder="Enter client and case info..."),
-        gr.File(label="Reference Template (DOCX/PDF/TXT)", file_types=[".docx",".pdf",".txt"]),
-        gr.Checkbox(label="Append citations", value=True)
-    ],
-    outputs=[
-        gr.Textbox(label="Draft Output", lines=30),
-        gr.File(label="Download Word")
-    ],
-    title="⚖️ AI Legal Draft Generator",
-    description="Upload a DOCX/PDF/TXT reference template, enter case details, and generate a court-ready legal draft in Markdown and Word."
-)
+with gr.Blocks() as iface:
+    gr.Markdown("## ⚖️ AI Legal Draft Generator\nUpload a DOCX/PDF/TXT template and enter case details. Get a court-ready draft.")
+
+    case_text = gr.Textbox(label="Case Details", lines=10, placeholder="Enter client and case info...")
+    uploaded_file = gr.File(label="Reference Template (DOCX/PDF/TXT)", file_types=[".docx",".pdf",".txt"])
+    add_citations = gr.Checkbox(label="Append citations", value=True)
+
+    draft_output = gr.Markdown(label="Draft Output")
+    download_btn = gr.DownloadButton(label="⬇️ Download Word", file_name="legal_draft.docx")
+
+    btn = gr.Button("Generate Draft")
+    btn.click(
+        generate_legal_draft,
+        inputs=[case_text, uploaded_file, add_citations],
+        outputs=[draft_output, download_btn],
+    )
 
 if __name__ == "__main__":
     iface.launch()