Update main.py

main.py

@@ -5,6 +5,7 @@ import math
 import asyncio
 from functools import lru_cache
 from typing import Any, Dict, List
+from utils import extract_json_obj_robust, sanitize_analyze_output
 from fastapi.middleware.cors import CORSMiddleware
 import uvicorn
 from fastapi import FastAPI
@@ -103,39 +104,35 @@ LABEL_MAPPING = {
     "intentional": ["intentional"]
 }
 
-ANALYZE_SYS_PROMPT = """You are a logic expert. Detect logical fallacies.
-OUTPUT JSON ONLY.
 
-RULES:
-1. Use ONLY these labels: {labels}
-2. "rationale": Explain WHY.
-3. "confidence": 0.0 to 1.0.
+ANALYZE_SYS_PROMPT = """You are a logic expert. Detect logical fallacies in the given text.
+OUTPUT JSON ONLY. No markdown. No extra keys. No commentary outside JSON.
 
-EXAMPLES (Follow this logic):
+IMPORTANT:
+- The text can contain MULTIPLE fallacies. Return ALL that apply.
+- If there are NO fallacies, set "has_fallacy": false and "fallacies": [].
+- "evidence_quotes" MUST be the SHORTEST exact span(s) from the input that justify the fallacy.
+  Do NOT quote the whole text. Prefer 1 short quote; at most 3 quotes.
 
-Input: "You are stupid, so your opinion is wrong."
-Output: {{
-  "has_fallacy": true,
-  "fallacies": [{{
-      "type": "ad hominem",
-      "confidence": 0.95,
-      "evidence_quotes": ["You are stupid"],
-      "rationale": "Direct attack on the person rather than the argument."
-  }}],
-  "overall_explanation": "Ad Hominem attack."
-}}
+LABELS:
+Use ONLY these labels: {labels}
+Do NOT invent labels.
+Do NOT output "none" as a fallacy item.
 
-Input: "Think of the children! We must ban this immediately or they will suffer!"
-Output: {{
-  "has_fallacy": true,
-  "fallacies": [{{
-      "type": "appeal to emotion",
-      "confidence": 0.90,
-      "evidence_quotes": ["Think of the children", "they will suffer"],
-      "rationale": "Uses fear and pity to manipulate opinion without logical proof."
-  }}],
-  "overall_explanation": "Manipulative emotional appeal."
-}}
+RATIONALE QUALITY (VERY IMPORTANT):
+Each fallacy "rationale" MUST be directly tied to THIS input text, and MUST NOT be generic.
+Structure each rationale like this (2–3 sentences max):
+1) Restate the specific claim from the input in your own words AND anchor it to the exact quote(s).
+2) Explain why that claim matches the fallacy, referencing what is missing or what is assumed.
+3) If relevant, mention a concrete cue in the text.
+
+OVERALL EXPLANATION (MUST reference the input):
+- First: a quick recap of the detected fallacies (1 short sentence).
+- Then: a general explanation of why these fallacies are risky IN THIS TEXT.
+- If no fallacy: briefly explain why the reasoning is acceptable / what would be needed to call it fallacious.
+
+CONFIDENCE:
+"confidence" is between 0.0 and 1.0.
 
 JSON SCHEMA:
 {{
@@ -150,40 +147,53 @@ JSON SCHEMA:
   ],
   "overall_explanation": string
 }}
+
+EXAMPLES (style guide — copy this style):
+
+Input: "If we allow remote work, productivity will collapse and the company will fail."
+Output:
+{{
+  "has_fallacy": true,
+  "fallacies": [{{
+    "type": "false causality",
+    "confidence": 0.86,
+    "evidence_quotes": ["If we allow remote work, productivity will collapse", "the company will fail"],
+    "rationale": "The input implies that allowing remote work will directly cause productivity to collapse and lead to company failure (quotes: 'productivity will collapse', 'the company will fail') without supporting evidence. It treats a speculative outcome as a guaranteed causal chain, jumping from a policy change to extreme consequences."
+  }}],
+  "overall_explanation": "Recap: false causality. Risk: the text presents a shaky cause-and-effect chain as certain, which can push decisions based on fear rather than evidence and ignore alternative explanations."
+}}
 """
 
-REWRITE_SYS_PROMPT = """You are a text editor. Rewrite to remove the fallacy.
-Output Format (JSON):
+REWRITE_SYS_PROMPT = """You are a careful editor. Rewrite the text to REMOVE the logical fallacy while PRESERVING the original meaning as much as possible.
+
+Context:
+- Predicted fallacy type: {fallacy_type}
+- Rationale: {rationale}
+
+GOAL:
+- Keep the same overall intent, but soften / qualify claims so the reasoning is no longer fallacious.
+- Avoid absolute language ("always", "everyone", "no one") unless fully justified in the text.
+- Replace overgeneralizations with reasonable qualifiers ("some", "often", "can", "in some cases").
+- If the issue is causality, add uncertainty or evidence requirements ("may contribute", "could be related", "without evidence we can't conclude").
+- If the issue is a false dilemma, add alternatives and nuance.
+- If the issue is ad hominem / credibility attacks, remove personal attacks and focus on claims/evidence.
+
+OUTPUT JSON ONLY (no markdown):
 {{
   "rewritten_text": string,
   "why_this_fix": string
 }}
-"""
 
-def clean_and_repair_json(text: str) -> str:
-    text = text.replace("```json", "").replace("```", "").strip()
-    start = text.find("{")
-    if start == -1:
-        return text
-
-    depth = 0
-    for i, char in enumerate(text[start:], start=start):
-        if char == "{":
-            depth += 1
-        elif char == "}":
-            depth -= 1
-            if depth == 0:
-                potential_json = text[start:i + 1]
-                try:
-                    json.loads(potential_json)
-                    return potential_json
-                except Exception:
-                    pass
-
-    end = text.rfind("}")
-    if start != -1 and end != -1:
-        return text[start:end + 1]
-    return text
+The "why_this_fix" must be short (1-2 sentences) and explain what you changed to remove the fallacy.
+
+EXAMPLE:
+Input idea: "All blond women are pretty."
+Output:
+{{
+  "rewritten_text": "Some blond women can be very pretty, but attractiveness varies from person to person.",
+  "why_this_fix": "It removes the absolute generalization and replaces it with a qualified statement that doesn't stereotype an entire group."
+}}
+"""
 
 def analyze_alternatives(start_index: int, top_logprobs_list: List[Dict[str, float]]) -> Dict[str, float]:
     if start_index < 0 or start_index >= len(top_logprobs_list):
@@ -336,14 +346,15 @@ async def analyze(req: AnalyzeRequest):
         logprobs = lp_data.get("token_logprobs", [])
         top_logprobs = lp_data.get("top_logprobs", [])
 
-    cleaned_text = clean_and_repair_json(raw_text)
+    parsed_obj = extract_json_obj_robust(raw_text)
     result_json: Dict[str, Any] = {}
     success = False
     technical_confidence = 0.0
     label_distribution: Dict[str, float] = {}
 
-    try:
-        result_json = json.loads(cleaned_text)
+    if parsed_obj is not None:
+        # Enforce schema + clean common template artifacts
+        result_json = sanitize_analyze_output(parsed_obj, req.text)
         success = True
 
         if result_json.get("has_fallacy") and result_json.get("fallacies"):
@@ -358,7 +369,7 @@ async def analyze(req: AnalyzeRequest):
                     fallacy["alternatives"] = label_distribution
 
                     declared = fallacy.get("confidence", 0.8)
-                    fallacy["confidence"] = round((declared + spec_conf) / 2, 2)
+                    fallacy["confidence"] = round((float(declared) + float(spec_conf)) / 2, 2)
 
                     if technical_confidence == 0.0:
                         technical_confidence = spec_conf
@@ -367,7 +378,7 @@ async def analyze(req: AnalyzeRequest):
                 info = extract_label_info("has_fallacy", tokens, logprobs, top_logprobs)
                 label_distribution = info["dist"]
 
-    except json.JSONDecodeError:
+    else:
         result_json = {"error": "JSON Error", "raw": raw_text}
         success = False
 
@@ -392,10 +403,13 @@ async def rewrite(req: RewriteRequest):
             max_tokens=req.max_new_tokens,
             temperature=0.7,
             repeat_penalty=1.1,
-            stop=["</s>", "}"],
+            stop=["</s>", "```"],
         )
     try:
-        res = json.loads(clean_and_repair_json(output["choices"][0]["text"]))
+        parsed = extract_json_obj_robust(output["choices"][0]["text"])
+        if parsed is None:
+            raise ValueError("json_parse_failed")
+        res = parsed
         ok = True
     except Exception:
         res = {"raw": output["choices"][0]["text"]}