Update app.py

app.py

@@ -61,15 +61,6 @@ def load_templates_json() -> List[Dict]:
     except Exception:
         return []
 
-# def parse_json_safe(raw_text: str, fallback: str) -> List[str]:
-#     try:
-#         return json.loads(raw_text)
-#     except:
-#         matches = re.findall(r'"([^"]+)"', raw_text)
-#         if matches:
-#             return matches
-#         return [fallback[:512]]
-
 def save_template_to_hf(name: str, analysis: str) -> Tuple[bool, str]:
     """Save new template into HF dataset repo (templates.json)."""
     try:
@@ -120,7 +111,7 @@ def parse_json_safe(raw_text: str, fallback: str) -> List[str]:
             return matches
         return [fallback[:512]]
 
-def build_queries_with_llm(user_text: str, max_queries: int = 20) -> List[str]:
+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. "
@@ -129,23 +120,18 @@ def build_queries_with_llm(user_text: str, max_queries: int = 20) -> List[str]:
         "and FBR tax ordinances. "
         "Return ONLY a JSON array of strings, no extra text."
     )
-    user_prompt = f"Case description:\n{user_text}"
-
     try:
         resp = openai_client.chat.completions.create(
             model="gpt-4o-mini",
-            messages=[
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": user_prompt}
-            ],
+            messages=[{"role": "system", "content": system_prompt},
+                      {"role": "user", "content": user_text}],
             temperature=0.1,
             max_tokens=2000
         )
         raw = resp.choices[0].message.content.strip()
-        queries = parse_json_safe(raw, user_text)
+        return parse_json_safe(raw, user_text)[:max_queries]
     except Exception:
-        queries = [user_text[:512]]
-    return queries[:max_queries]
+        return [user_text[:512]]
 
 def pinecone_search(queries: List[str], top_k: int = 10, max_chars: int = 10000) -> Tuple[str, List[Dict]]:
     seen_texts, context_parts, citations = set(), [], []
@@ -169,10 +155,8 @@ def pinecone_search(queries: List[str], top_k: int = 10, max_chars: int = 10000)
     return "\n".join(context_parts), citations
 
 def markdown_to_docx(md_text: str) -> str:
-    """Convert Markdown text into a Word DOCX and return a file path."""
     html = markdown.markdown(md_text)
     soup = BeautifulSoup(html, "html.parser")
-
     doc = Document()
     for el in soup.descendants:
         if el.name == "h1":
@@ -185,17 +169,14 @@ def markdown_to_docx(md_text: str) -> str:
             doc.add_paragraph(el.get_text())
         elif el.name == "li":
             doc.add_paragraph(f"• {el.get_text()}")
-
     tmp_path = os.path.join(tempfile.gettempdir(), "draft.docx")
     doc.save(tmp_path)
     return tmp_path
 
-# ----------------- DRAFT ANALYZER -----------------
+# ----------------- 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:
@@ -207,32 +188,26 @@ Extract the following information clearly and systematically:
 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
+            model="gpt-5",
+            messages=[{"role": "system", "content": system_prompt},
+                      {"role": "user", "content": ref_text[:40000]}],
+            max_completion_tokens=4000,
+            verbosity="medium"
         )
-        analysis = resp.choices[0].message.content.strip()
+        return resp.choices[0].message.content.strip()
     except Exception as e:
-        analysis = f"(Analyzer error: {e})"
-
-    return analysis
+        return f"(Analyzer error: {e})"
 
 # ----------------- MAIN -----------------
 def generate_legal_draft(case_text, uploaded_file, template_name, new_template_name, add_citations=True):
     yield gr.update(value="🔍 Searching in Knowledge Base..."), None
 
     queries = build_queries_with_llm(case_text)
-    context_text, citations = pinecone_search(queries, top_k=15)
+    context_text, citations = pinecone_search(queries, top_k=10)
 
     # Handle template
     template_analysis = ""
@@ -269,10 +244,8 @@ def generate_legal_draft(case_text, uploaded_file, template_name, new_template_n
     user_prompt = f"""
 **User Input:** 
 {case_text}
-
 **Knowledge Base Context:**
 {context_text or '(no matches)'}
-
 **Template Draft Analysis:**
 {template_analysis}
 """
@@ -280,10 +253,8 @@ def generate_legal_draft(case_text, uploaded_file, template_name, new_template_n
     try:
         resp = openai_client.chat.completions.create(
             model="gpt-5",
-            messages=[
-                {"role":"system","content":system_prompt},
-                {"role":"user","content":user_prompt}
-            ],
+            messages=[{"role": "system", "content": system_prompt},
+                      {"role": "user", "content": user_prompt}],
             max_completion_tokens=15000,
             verbosity="high"
         )
@@ -297,9 +268,7 @@ def generate_legal_draft(case_text, uploaded_file, template_name, new_template_n
             draft_md += f"{i}. {c['source']} (score: {c['score']:.3f})\n"
 
     docx_path = markdown_to_docx(draft_md)
-    yield gr.update(value=draft_md), markdown_to_docx(draft_md)
-
-
+    yield gr.update(value=draft_md), docx_path
 
 # ----------------- GRADIO UI -----------------
 with gr.Blocks() as demo:
@@ -320,9 +289,9 @@ with gr.Blocks() as demo:
 
     btn.click(
         generate_legal_draft,
-        inputs=[case_text, template_name,uploaded_file, new_template_name, add_citations],
+        inputs=[case_text, uploaded_file, template_name, new_template_name, add_citations],
         outputs=[draft_output, download_btn]
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()