Update ui/ui_core.py

ui/ui_core.py

@@ -16,31 +16,13 @@ from txagent.txagent import TxAgent
 def sanitize_utf8(text: str) -> str:
     return text.encode("utf-8", "ignore").decode("utf-8")
 
-def clean_final_response(text: str) -> str:
-    cleaned = text.replace("[TOOL_CALLS]", "").strip()
-    responses = cleaned.split("[Final Analysis]")
-
-    if len(responses) <= 1:
-        return f"<div style='padding:1em;border:1px solid #ccc;border-radius:12px;color:#fff;background:#353F54;'><p>{cleaned}</p></div>"
-
-    panels = []
-    for i, section in enumerate(responses[1:], 1):
-        final = section.strip()
-        panels.append(
-            f"<div style='background:#2B2B2B;color:#E0E0E0;border-radius:12px;margin-bottom:1em;border:1px solid #888;'>"
-            f"<div style='font-size:1.1em;font-weight:bold;padding:0.75em;background:#3A3A3A;color:#fff;border-radius:12px 12px 0 0;'>🧠 Final Analysis #{i}</div>"
-            f"<div style='padding:1em;line-height:1.6;'>{final.replace(chr(10), '<br>')}</div>"
-            f"</div>"
-        )
-    return "".join(panels)
-
 def file_hash(path):
     with open(path, "rb") as f:
         return hashlib.md5(f.read()).hexdigest()
 
 def convert_file_to_json(file_path: str, file_type: str) -> str:
     try:
-        cache_dir = os.path.join("cache")
+        cache_dir = "/data/cache"
         os.makedirs(cache_dir, exist_ok=True)
         h = file_hash(file_path)
         cache_path = os.path.join(cache_dir, f"{h}.json")
@@ -75,23 +57,6 @@ def convert_file_to_json(file_path: str, file_type: str) -> str:
     except Exception as e:
         return json.dumps({"error": f"Error reading {os.path.basename(file_path)}: {str(e)}"})
 
-def chunk_text(text: str, max_tokens: int = 6000) -> List[str]:
-    chunks = []
-    words = text.split()
-    chunk = []
-    token_count = 0
-    for word in words:
-        token_count += len(word) // 4 + 1
-        if token_count > max_tokens:
-            chunks.append(" ".join(chunk))
-            chunk = [word]
-            token_count = len(word) // 4 + 1
-        else:
-            chunk.append(word)
-    if chunk:
-        chunks.append(" ".join(chunk))
-    return chunks
-
 def create_ui(agent: TxAgent):
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
         gr.Markdown("<h1 style='text-align: center;'>📋 CPS: Clinical Patient Support System</h1>")
@@ -107,15 +72,6 @@ def create_ui(agent: TxAgent):
         conversation_state = gr.State([])
 
         def handle_chat(message: str, history: list, conversation: list, uploaded_files: list, progress=gr.Progress()):
-            context = (
-                "You are an expert clinical AI assistant reviewing medical form or interview data. "
-                "Your job is to analyze this data and reason about any information or red flags that a human doctor might have overlooked. "
-                "Provide a **detailed and structured response**, including examples, supporting evidence from the form, and clinical rationale for why these items matter. "
-                "Ensure the output is informative and helpful for improving patient care. "
-                "Do not hallucinate. Base the response only on the provided form content. "
-                "End with a section labeled '[Final Analysis]' where you summarize key findings the doctor may have missed."
-            )
-
             try:
                 history.append({"role": "user", "content": message})
                 history.append({"role": "assistant", "content": "⏳ Processing your request..."})
@@ -131,47 +87,37 @@ def create_ui(agent: TxAgent):
                         json_text = convert_file_to_json(path, ext)
                         extracted_text += sanitize_utf8(json_text) + "\n"
 
-                chunks = chunk_text(extracted_text.strip())
-
-                def process_chunk(i, chunk):
-                    if len(chunk.split()) > 8192:
-                        return f"[Chunk {i+1}] ⚠️ Skipped: input exceeds model limit."
-
-                    chunked_prompt = (
-                        f"{context}\n\n--- Uploaded File Content (Chunk {i+1}/{len(chunks)}) ---\n\n{chunk}\n\n"
-                        f"--- End of Chunk ---\n\nNow begin your analysis:"
-                    )
-                    try:
-                        generator = agent.run_gradio_chat(
-                            message=chunked_prompt,
-                            history=[],
-                            temperature=0.3,
-                            max_new_tokens=1024,
-                            max_token=8192,
-                            call_agent=False,
-                            conversation=conversation,
-                            uploaded_files=uploaded_files,
-                            max_round=30
-                        )
-                        result = ""
-                        for update in generator:
-                            if update is None:
-                                continue
-                            if isinstance(update, str):
-                                result += update
-                            elif isinstance(update, list):
-                                for msg in update:
-                                    if hasattr(msg, 'content'):
-                                        result += msg.content
-                        return result if result.strip() else f"[Chunk {i+1}] ⚠️ No response received."
-                    except Exception as err:
-                        print(f"[Error in chunk {i+1}] {err}")
-                        return f"[Chunk {i+1}] ❌ Failed to process due to error."
-
-                results = [process_chunk(i, chunk) for i, chunk in enumerate(chunks)]
-                full_response = "\n\n".join(results)
-                full_response = clean_final_response(full_response.strip())
-                history[-1] = {"role": "assistant", "content": full_response}
+                # Only final chunk will be passed (no split or loop)
+                context = (
+                    "You are an expert clinical AI assistant. Review this patient's history, medications, and notes, and ONLY provide a final answer summarizing what the doctor might have missed."
+                )
+
+                chunked_prompt = f"{context}\n\n--- Patient Record ---\n{extracted_text}\n\n[Final Analysis]"
+
+                generator = agent.run_gradio_chat(
+                    message=chunked_prompt,
+                    history=[],
+                    temperature=0.3,
+                    max_new_tokens=1024,
+                    max_token=8192,
+                    call_agent=False,
+                    conversation=conversation,
+                    uploaded_files=uploaded_files,
+                    max_round=30
+                )
+
+                final_response = ""
+                for update in generator:
+                    if update is None:
+                        continue
+                    if isinstance(update, str):
+                        final_response += update
+                    elif isinstance(update, list):
+                        for msg in update:
+                            if hasattr(msg, 'content'):
+                                final_response += msg.content
+
+                history[-1] = {"role": "assistant", "content": final_response.strip() or "❌ No response."}
                 yield history
 
             except Exception as chat_error:
@@ -189,4 +135,4 @@ def create_ui(agent: TxAgent):
             ["Is there anything abnormal in the attached blood work report?"]
         ], inputs=message_input)
 
-    return demo
+    return demo