Update app.py

app.py

@@ -1,4 +1,4 @@
-# (Full Updated Code Snippet with Proper Final Response Handling)
+# Full updated app.py with TOOL_CALLS displayed separately and full fixes
 
 import sys
 import os
@@ -23,6 +23,7 @@ if not torch.cuda.is_available():
 
 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")
@@ -67,10 +68,9 @@ def extract_priority_pages(file_path: str, max_pages: int = 20) -> str:
                     text_chunks.append(f"=== Page {i} ===\n{page_text.strip()}")
         return "\n\n".join(text_chunks)
     except Exception as e:
-        debug_msg = f"PDF processing error: {str(e)}"
-        print(debug_msg)
+        print("PDF processing error:", str(e))
         traceback.print_exc()
-        return debug_msg
+        return str(e)
 
 def convert_file_to_json(file_path: str, file_type: str) -> str:
     try:
@@ -84,8 +84,7 @@ def convert_file_to_json(file_path: str, file_type: str) -> str:
             text = extract_priority_pages(file_path)
             result = json.dumps({"filename": os.path.basename(file_path), "content": text, "status": "initial"})
         elif file_type == "csv":
-            df = pd.read_csv(file_path, encoding_errors="replace", header=None, dtype=str,
-                             skip_blank_lines=False, on_bad_lines="skip")
+            df = pd.read_csv(file_path, encoding_errors="replace", header=None, dtype=str, skip_blank_lines=False, on_bad_lines="skip")
             content = df.fillna("").astype(str).values.tolist()
             result = json.dumps({"filename": os.path.basename(file_path), "rows": content})
         elif file_type in ["xls", "xlsx"]:
@@ -97,14 +96,14 @@ def convert_file_to_json(file_path: str, file_type: str) -> str:
             result = json.dumps({"filename": os.path.basename(file_path), "rows": content})
         else:
             result = json.dumps({"error": f"Unsupported file type: {file_type}"})
+
         with open(cache_path, "w", encoding="utf-8") as f:
             f.write(result)
         return result
     except Exception as e:
-        error_msg = f"Error processing {os.path.basename(file_path)}: {str(e)}"
-        print(error_msg)
+        print("Error processing", file_path, str(e))
         traceback.print_exc()
-        return json.dumps({"error": error_msg})
+        return json.dumps({"error": str(e)})
 
 def log_system_usage(tag=""):
     try:
@@ -124,7 +123,7 @@ def log_system_usage(tag=""):
 
 def init_agent():
     try:
-        print("\U0001F501 Initializing model...")
+        print("🔁 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")
@@ -173,14 +172,14 @@ def create_ui(agent):
                         results = []
                         for future in as_completed(futures):
                             try:
-                                results.append(sanitize_utf8(future.result()))
+                                res = future.result()
+                                results.append(sanitize_utf8(res))
                             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)
+                    extracted = "\n".join(results)
+                    file_hash_value = file_hash(files[0].name)
 
-                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
@@ -188,13 +187,16 @@ def create_ui(agent):
 4. Highlight abnormal results needing follow-up
 
 Medical Records:
-{extracted[:max_content_length]}
+{extracted[:8000]}
 
 ### Potential Oversights:
 """
+                print("🔎 Generated prompt:")
+                print(prompt)
 
                 full_response = ""
                 response_chunks = []
+                tool_calls_rendered = []
 
                 for chunk in agent.run_gradio_chat(
                     message=prompt,
@@ -205,55 +207,39 @@ Medical Records:
                     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)
-
-                        display_response = re.split(r"\\[TOOL_CALLS\\].*?$", full_response, flags=re.DOTALL)[0].strip()
-                        display_response = display_response.replace('[TxAgent]', '').strip()
-
-                        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()
+                    if chunk is None:
+                        continue
+                    chunk_content = chunk if isinstance(chunk, str) else getattr(chunk, 'content', '')
+                    if not chunk_content:
                         continue
+                    response_chunks.append(chunk_content)
+                    full_response = "".join(response_chunks)
 
-                if not full_response:
-                    full_response = "⚠️ No clear oversights identified or model output was invalid."
-                else:
-                    full_response = re.split(r"\\[TOOL_CALLS\\].*?$", full_response, flags=re.DOTALL)[0].strip()
-                    full_response = full_response.replace('[TxAgent]', '').strip()
+                    matches = re.findall(r"\[TOOL_CALLS\]\[(.*?)\]", chunk_content, re.DOTALL)
+                    for m in matches:
+                        tool_calls_rendered.append(f"\n📦 Tool Call: [{m.strip()}]")
+
+                    display_response = re.sub(r"\[TOOL_CALLS\].*?\n*", "", full_response, flags=re.DOTALL)
+                    display_response = display_response.replace('[TxAgent]', '').strip()
+                    display_response += "\n\n" + "\n".join(tool_calls_rendered)
+
+                    if history and history[-1]["role"] == "assistant":
+                        history[-1]["content"] = display_response
+                    else:
+                        history.append({"role": "assistant", "content": display_response})
+
+                    yield history, None
+
+                full_response = re.sub(r"\[TOOL_CALLS\].*?\n*", "", full_response, flags=re.DOTALL).strip()
+                full_response = full_response.replace('[TxAgent]', '').strip()
 
                 report_path = None
                 if file_hash_value:
                     report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt")
-                    try:
-                        with open(report_path, "w", encoding="utf-8") as f:
-                            f.write(full_response)
-                    except Exception as e:
-                        print("❌ Error saving report:", str(e))
-                        traceback.print_exc()
-
-                if len(history) > 0 and history[-1]["role"] == "assistant":
+                    with open(report_path, "w", encoding="utf-8") as f:
+                        f.write(full_response)
+
+                if history and history[-1]["role"] == "assistant":
                     history[-1]["content"] = full_response
                 else:
                     history.append({"role": "assistant", "content": full_response})
@@ -261,10 +247,8 @@ Medical Records:
                 yield history, report_path if report_path and os.path.exists(report_path) else None
 
             except Exception as e:
-                error_message = f"❌ An error occurred in analyze: {str(e)}"
-                print(error_message)
+                history.append({"role": "assistant", "content": f"❌ An error occurred in analyze: {str(e)}"})
                 traceback.print_exc()
-                history.append({"role": "assistant", "content": error_message})
                 yield history, None
 
         send_btn.click(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output])