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,6 +108,17 @@ 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")
@@ -131,33 +142,6 @@ def init_agent():
     print("✅ Agent Ready")
     return agent
 
-def clean_response(response: str) -> str:
-    """Clean the response by removing tool calls and duplicate content."""
-    # First remove all tool call blocks
-    response = re.sub(r'\[TOOL_CALLS\].*?(\[TOOL_CALLS\]|$)', '', response, flags=re.DOTALL)
-    
-    # Then remove any remaining standalone tool call markers
-    response = response.replace('[TOOL_CALLS]', '')
-    
-    # Remove duplicate sections (looking for repeated identical paragraphs)
-    paragraphs = [p.strip() for p in response.split('\n\n') if p.strip()]
-    unique_paragraphs = []
-    seen_paragraphs = set()
-    
-    for para in paragraphs:
-        if para not in seen_paragraphs:
-            seen_paragraphs.add(para)
-            unique_paragraphs.append(para)
-    
-    # Reconstruct the response
-    cleaned = '\n\n'.join(unique_paragraphs)
-    
-    # Remove any remaining JSON-like artifacts
-    cleaned = re.sub(r'\{.*?\}', '', cleaned)
-    cleaned = re.sub(r'\[.*?\]', '', cleaned)
-    
-    return cleaned.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>")
@@ -194,7 +178,7 @@ Medical Records:
 """
 
             try:
-                full_response = ""
+                response = ""
                 for chunk in agent.run_gradio_chat(
                     message=prompt,
                     history=[],
@@ -207,29 +191,19 @@ Medical Records:
                     if chunk is None:
                         continue
                     if isinstance(chunk, str):
-                        full_response += chunk
+                        response += chunk
                     elif isinstance(chunk, list):
-                        full_response += "".join([c.content for c in chunk if hasattr(c, "content") and c.content])
-
-                    # Clean the current response for display
-                    current_cleaned = clean_response(full_response)
-                    if current_cleaned:
-                        history[-1] = {"role": "assistant", "content": current_cleaned}
-                        yield history, None
-
-                # Final cleaning and processing
-                final_cleaned = clean_response(full_response)
-                if not final_cleaned:
-                    final_cleaned = "⚠️ No clear oversights identified or model output was invalid."
-
-                # Save the full report
-                report_path = None
-                if 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}
+                        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
                 yield history, report_path if report_path and os.path.exists(report_path) else None
 
             except Exception as e: