Update app.py

app.py

@@ -11,6 +11,10 @@ import tempfile
 import json
 import re
 
+import markdown
+from bs4 import BeautifulSoup
+
+
 # ----------------- CONFIG -----------------
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
@@ -69,7 +73,7 @@ def build_queries_with_llm(user_text: str, max_queries: int = 15) -> List[str]:
                 {"role": "system", "content": system_prompt},
                 {"role": "user", "content": user_prompt}
             ],
-            temperature=0.2,
+            temperature=0.1,
             max_tokens=2000
         )
         raw = resp.choices[0].message.content.strip()
@@ -99,21 +103,45 @@ def pinecone_search(queries: List[str], top_k: int = 10, max_chars: int = 10000)
                 break
     return "\n".join(context_parts), citations
 
-def markdown_to_docx(md_text: str) -> str:
-    """
-    Convert Markdown text into a temporary .docx file and return its path.
-    """
-    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".docx")
+# def markdown_to_docx(md_text: str) -> str:
+#     """
+#     Convert Markdown text into a temporary .docx file and return its path.
+#     """
+#     temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".docx")
+#     doc = Document()
+#     for line in md_text.split("\n"):
+#         if line.startswith("## "):
+#             doc.add_heading(line[3:], level=2)
+#         elif line.startswith("# "):
+#             doc.add_heading(line[2:], level=1)
+#         else:
+#             doc.add_paragraph(line)
+#     doc.save(temp_file.name)
+#     return temp_file.name
+    
+def markdown_to_docx(md_text: str) -> BytesIO:
+    """Convert Markdown text into Word DOCX with proper formatting."""
+    html = markdown.markdown(md_text)
+    soup = BeautifulSoup(html, "html.parser")
+
     doc = Document()
-    for line in md_text.split("\n"):
-        if line.startswith("## "):
-            doc.add_heading(line[3:], level=2)
-        elif line.startswith("# "):
-            doc.add_heading(line[2:], level=1)
-        else:
-            doc.add_paragraph(line)
-    doc.save(temp_file.name)
-    return temp_file.name
+    for el in soup.descendants:
+        if el.name == "h1":
+            doc.add_heading(el.get_text(), level=1)
+        elif el.name == "h2":
+            doc.add_heading(el.get_text(), level=2)
+        elif el.name == "h3":
+            doc.add_heading(el.get_text(), level=3)
+        elif el.name == "p":
+            doc.add_paragraph(el.get_text())
+        elif el.name == "li":
+            doc.add_paragraph(f"• {el.get_text()}")
+
+    buf = BytesIO()
+    doc.save(buf)
+    buf.seek(0)
+    return buf
+
 
 # ----------------- MAIN FUNCTION -----------------
 def generate_legal_draft(case_text, uploaded_file=None, add_citations=True):
@@ -168,7 +196,7 @@ def generate_legal_draft(case_text, uploaded_file=None, add_citations=True):
                 {"role":"system","content":system_prompt},
                 {"role":"user","content":user_prompt}
             ],
-            temperature=0.2,
+            temperature=0.1,
             max_tokens=15000
         )
         draft_md = resp.choices[0].message.content.strip()
@@ -181,7 +209,8 @@ 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 draft_md, docx_path
+    return gr.update(value=draft_md), markdown_to_docx(draft_md)
+
 
 # ----------------- GRADIO INTERFACE -----------------
 with gr.Blocks() as demo:
@@ -192,6 +221,8 @@ with gr.Blocks() as demo:
     add_citations = gr.Checkbox(label="Append citations", value=True)
 
     draft_output = gr.Markdown(label="Draft Output")
+    draft_md = draft_md.replace("\\n", "\n")  # convert escaped newlines
+
     download_btn = gr.DownloadButton(label="⬇️ Download Word")
     # draft_output1 = gr.Markdown("# HI")
 
@@ -200,7 +231,8 @@ with gr.Blocks() as demo:
         generate_legal_draft,
         inputs=[case_text, uploaded_file, add_citations],
         outputs=[draft_output, download_btn],
-    )
+)
+
 
 # ----------------- RUN -----------------
 if __name__ == "__main__":