Update app.py

app.py

@@ -1,3 +1,5 @@
+# (Full Updated Code Snippet with Proper Final Response Handling)
+
 import sys
 import os
 import pandas as pd
@@ -14,18 +16,13 @@ import subprocess
 import traceback
 import torch
 
-# Set VLLM logging level to DEBUG for detailed output
 os.environ["VLLM_LOGGING_LEVEL"] = "DEBUG"
-
-# If no GPU is available, force CPU usage by hiding CUDA devices
 if not torch.cuda.is_available():
     print("No GPU detected. Forcing CPU mode by setting CUDA_VISIBLE_DEVICES to an empty string.")
     os.environ["CUDA_VISIBLE_DEVICES"] = ""
 
-# Persistent directory setup
 persistent_dir = "/data/hf_cache"
 os.makedirs(persistent_dir, exist_ok=True)
-
 model_cache_dir = os.path.join(persistent_dir, "txagent_models")
 tool_cache_dir = os.path.join(persistent_dir, "tool_cache")
 file_cache_dir = os.path.join(persistent_dir, "cache")
@@ -35,7 +32,6 @@ vllm_cache_dir = os.path.join(persistent_dir, "vllm_cache")
 for directory in [model_cache_dir, tool_cache_dir, file_cache_dir, report_dir, vllm_cache_dir]:
     os.makedirs(directory, exist_ok=True)
 
-# Update environment variables
 os.environ["HF_HOME"] = model_cache_dir
 os.environ["TRANSFORMERS_CACHE"] = model_cache_dir
 os.environ["VLLM_CACHE_DIR"] = vllm_cache_dir
@@ -48,7 +44,6 @@ sys.path.insert(0, src_path)
 
 from txagent.txagent import TxAgent
 
-# Medical keywords for processing PDF files
 MEDICAL_KEYWORDS = {'diagnosis', 'assessment', 'plan', 'results', 'medications',
                     'allergies', 'summary', 'impression', 'findings', 'recommendations'}
 
@@ -68,7 +63,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:
@@ -129,7 +124,7 @@ def log_system_usage(tag=""):
 
 def init_agent():
     try:
-        print("🔁 Initializing model...")
+        print("\U0001F501 Initializing model...")
         log_system_usage("Before Load")
         default_tool_path = os.path.abspath("data/new_tool.json")
         target_tool_path = os.path.join(tool_cache_dir, "new_tool.json")
@@ -166,33 +161,26 @@ def create_ui(agent):
 
         def analyze(message: str, history: list, files: list):
             try:
-                # Initialize response with loading message
                 history.append({"role": "user", "content": message})
                 history.append({"role": "assistant", "content": "⏳ Analyzing records for potential oversights..."})
                 yield history, None
 
-                # Process files in parallel
                 extracted = ""
                 file_hash_value = ""
                 if files:
                     with ThreadPoolExecutor(max_workers=4) as executor:
-                        futures = [
-                            executor.submit(convert_file_to_json, f.name, f.name.split(".")[-1].lower())
-                            for f in files
-                        ]
+                        futures = [executor.submit(convert_file_to_json, f.name, f.name.split(".")[-1].lower()) for f in files]
                         results = []
                         for future in as_completed(futures):
                             try:
-                                res = future.result()
-                                results.append(sanitize_utf8(res))
+                                results.append(sanitize_utf8(future.result()))
                             except Exception as e:
                                 print("❌ Error in file processing:", str(e))
                                 traceback.print_exc()
                         extracted = "\n".join(results)
                         file_hash_value = file_hash(files[0].name)
 
-                # Truncate extracted content to avoid token limit issues
-                max_content_length = 8000  # Reduced from 12000 to prevent token overflow
+                max_content_length = 8000
                 prompt = f"""Review these medical records and identify EXACTLY what might have been missed:
 1. List potential missed diagnoses
 2. Flag any medication conflicts
@@ -205,78 +193,56 @@ Medical Records:
 ### Potential Oversights:
 """
 
-                print("🔎 Generated prompt:")
-                print(prompt)
-
-                # Initialize response tracking
                 full_response = ""
                 response_chunks = []
 
-                # Process streaming response with error handling
-                try:
-                    for chunk in agent.run_gradio_chat(
-                        message=prompt,
-                        history=[],
-                        temperature=0.2,
-                        max_new_tokens=2048,
-                        max_token=4096,
-                        call_agent=False,
-                        conversation=[]
-                    ):
-                        try:
-                            if chunk is None:
-                                continue
-                            
-                            # Handle different chunk types
-                            if isinstance(chunk, str):
-                                chunk_content = chunk
-                            elif hasattr(chunk, 'content'):
-                                chunk_content = chunk.content
-                            elif isinstance(chunk, list):
-                                chunk_content = "".join([c.content for c in chunk if hasattr(c, "content") and c.content])
-                            else:
-                                print("DEBUG: Received unknown type chunk", type(chunk))
-                                continue
-
-                            if not chunk_content:
-                                continue
+                for chunk in agent.run_gradio_chat(
+                    message=prompt,
+                    history=[],
+                    temperature=0.2,
+                    max_new_tokens=2048,
+                    max_token=4096,
+                    call_agent=False,
+                    conversation=[]
+                ):
+                    try:
+                        chunk_content = ""
+                        if isinstance(chunk, str):
+                            chunk_content = chunk
+                        elif hasattr(chunk, 'content'):
+                            chunk_content = chunk.content
+                        elif isinstance(chunk, list):
+                            chunk_content = "".join([c.content for c in chunk if hasattr(c, "content") and c.content])
+
+                        if not chunk_content:
+                            continue
 
-                            response_chunks.append(chunk_content)
-                            full_response = "".join(response_chunks)
+                        response_chunks.append(chunk_content)
+                        full_response = "".join(response_chunks)
 
-                            # Clean the response for display
-                            display_response = full_response.split('[TOOL_CALLS]')[0].strip()
-                            display_response = display_response.replace('[TxAgent]', '').strip()
+                        display_response = re.split(r"\\[TOOL_CALLS\\].*?$", full_response, flags=re.DOTALL)[0].strip()
+                        display_response = display_response.replace('[TxAgent]', '').strip()
 
-                            # Update the chat history with the latest response
+                        if len(history) > 1 and history[-2]["role"] == "assistant" and history[-2]["content"] == display_response:
+                            pass
+                        else:
                             if len(history) > 0 and history[-1]["role"] == "assistant":
                                 history[-1]["content"] = display_response
                             else:
                                 history.append({"role": "assistant", "content": display_response})
 
-                            yield history, None
-
-                        except Exception as e:
-                            print("❌ Error processing chunk:", str(e))
-                            traceback.print_exc()
-                            continue
-
-                except Exception as e:
-                    print("❌ Error in model streaming:", str(e))
-                    traceback.print_exc()
-                    history.append({"role": "assistant", "content": f"Error in model response: {str(e)}"})
-                    yield history, None
-                    return
+                        yield history, None
+                    except Exception as e:
+                        print("❌ Error processing chunk:", str(e))
+                        traceback.print_exc()
+                        continue
 
-                # Final response handling
                 if not full_response:
                     full_response = "⚠️ No clear oversights identified or model output was invalid."
                 else:
-                    # Clean up the final response
-                    full_response = full_response.split('[TOOL_CALLS]')[0].strip()
+                    full_response = re.split(r"\\[TOOL_CALLS\\].*?$", full_response, flags=re.DOTALL)[0].strip()
                     full_response = full_response.replace('[TxAgent]', '').strip()
 
-                # Save report if we have files
                 report_path = None
                 if file_hash_value:
                     report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt")
@@ -287,7 +253,6 @@ Medical Records:
                         print("❌ Error saving report:", str(e))
                         traceback.print_exc()
 
-                # Ensure the final response is in the history
                 if len(history) > 0 and history[-1]["role"] == "assistant":
                     history[-1]["content"] = full_response
                 else: