Update app.py

app.py

@@ -56,7 +56,7 @@ def extract_priority_pages(file_path: str, max_pages: int = 20) -> str:
                 text_chunks.append(f"=== Page {i+1} ===\n{text.strip()}")
             for i, page in enumerate(pdf.pages[3:max_pages], start=4):
                 page_text = page.extract_text() or ""
-                if any(re.search(rf'\\b{kw}\\b', page_text.lower()) for kw in MEDICAL_KEYWORDS):
+                if any(re.search(rf'\b{kw}\b', page_text.lower()) for kw in MEDICAL_KEYWORDS):
                     text_chunks.append(f"=== Page {i} ===\n{page_text.strip()}")
         return "\n\n".join(text_chunks)
     except Exception as e:
@@ -108,17 +108,6 @@ def log_system_usage(tag=""):
     except Exception as e:
         print(f"[{tag}] GPU/CPU monitor failed: {e}")
 
-def extract_final_response(response: str) -> str:
-    try:
-        parts = response.split("[TOOL_CALLS]")
-        for i in reversed(range(len(parts))):
-            if i + 1 < len(parts) and '"name": "Finish"' in parts[i + 1]:
-                return parts[i].strip()
-        return response.strip()
-    except Exception as e:
-        print("❌ Failed to extract clean response:", str(e))
-        return response.strip()
-
 def init_agent():
     print("🔁 Initializing model...")
     log_system_usage("Before Load")
@@ -142,6 +131,22 @@ def init_agent():
     print("✅ Agent Ready")
     return agent
 
+def clean_response(response: str) -> str:
+    """Clean the response by removing tool calls and duplicate content."""
+    # Remove all tool call blocks
+    response = re.sub(r'\[TOOL_CALLS\].*?$', '', response, flags=re.DOTALL)
+    
+    # Remove duplicate content by keeping only the last occurrence
+    if "Based on the medical records provided" in response:
+        last_occurrence = response.rfind("Based on the medical records provided")
+        response = response[last_occurrence:]
+    
+    # Remove any remaining JSON artifacts
+    response = re.sub(r'\{.*?\}', '', response)
+    response = re.sub(r'\[.*?\]', '', response)
+    
+    return response.strip()
+
 def create_ui(agent):
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
         gr.Markdown("<h1 style='text-align: center;'>🩺 Clinical Oversight Assistant</h1>")
@@ -178,7 +183,9 @@ Medical Records:
 """
 
             try:
-                response = ""
+                full_response = ""
+                finish_detected = False
+                
                 for chunk in agent.run_gradio_chat(
                     message=prompt,
                     history=[],
@@ -190,20 +197,35 @@ Medical Records:
                 ):
                     if chunk is None:
                         continue
+                        
                     if isinstance(chunk, str):
-                        response += chunk
+                        full_response += chunk
                     elif isinstance(chunk, list):
-                        response += "".join([c.content for c in chunk if hasattr(c, "content") and c.content])
-
-                clean = extract_final_response(response)
-                print("🧼 Raw Response:\n", response)
-                print("✅ Cleaned Final Response:\n", clean)
-
-                if not clean:
-                    clean = "⚠️ No clear oversights identified or model output was invalid."
-
-                history[-1] = {"role": "assistant", "content": clean}
-                report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt") if file_hash_value else None
+                        chunk_content = "".join([c.content for c in chunk if hasattr(c, "content") and c.content])
+                        full_response += chunk_content
+                        if '"name": "Finish"' in chunk_content:
+                            finish_detected = True
+
+                    # Display intermediate response
+                    clean_intermediate = clean_response(full_response)
+                    if clean_intermediate:
+                        history[-1] = {"role": "assistant", "content": clean_intermediate}
+                        yield history, None
+
+                # Final processing after all chunks are received
+                final_cleaned = clean_response(full_response)
+                
+                if not final_cleaned:
+                    final_cleaned = "⚠️ No clear oversights identified or model output was invalid."
+                
+                # Only save report if we got a proper finish signal
+                report_path = None
+                if finish_detected and file_hash_value:
+                    report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt")
+                    with open(report_path, "w", encoding="utf-8") as f:
+                        f.write(final_cleaned)
+
+                history[-1] = {"role": "assistant", "content": final_cleaned}
                 yield history, report_path if report_path and os.path.exists(report_path) else None
 
             except Exception as e: