Update app.py

app.py

@@ -12,7 +12,7 @@ import re
 import psutil
 import subprocess
 import traceback
-import torch  # For checking CUDA availability
+import torch
 
 # Set VLLM logging level to DEBUG for detailed output
 os.environ["VLLM_LOGGING_LEVEL"] = "DEBUG"
@@ -35,7 +35,7 @@ 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 to use HF_HOME
+# 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
@@ -208,7 +208,6 @@ Medical Records:
 
                 # Initialize response tracking
                 full_response = ""
-                last_update_time = 0
                 response_chunks = []
 
                 # Process streaming response
@@ -224,10 +223,12 @@ Medical Records:
                     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:
@@ -240,11 +241,14 @@ Medical Records:
                         response_chunks.append(chunk_content)
                         full_response = "".join(response_chunks)
 
+                        # Remove any tool call markers from the displayed response
+                        display_response = full_response.split('[TOOL_CALLS]')[0].strip()
+
                         # Update the chat history with the latest response
                         if len(history) > 0 and history[-1]["role"] == "assistant":
-                            history[-1]["content"] = full_response
+                            history[-1]["content"] = display_response
                         else:
-                            history.append({"role": "assistant", "content": full_response})
+                            history.append({"role": "assistant", "content": display_response})
 
                         yield history, None
 
@@ -256,6 +260,9 @@ Medical Records:
                 # 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()
 
                 # Save report if we have files
                 report_path = None