Update app.py

app.py

@@ -32,6 +32,7 @@ sys.path.insert(0, src_path)
 
 from txagent.txagent import TxAgent
 
+# Constants
 MAX_MODEL_TOKENS = 32768
 MAX_CHUNK_TOKENS = 8192
 MAX_NEW_TOKENS = 2048
@@ -68,26 +69,19 @@ def split_text_into_chunks(text: str, max_tokens: int = MAX_CHUNK_TOKENS) -> Lis
     effective_max_tokens = max_tokens - PROMPT_OVERHEAD
     if effective_max_tokens <= 0:
         raise ValueError(f"Effective max tokens ({effective_max_tokens}) must be positive.")
-
     lines = text.split("\n")
-    chunks = []
-    current_chunk = []
-    current_tokens = 0
-
+    chunks, current_chunk, current_tokens = [], [], 0
     for line in lines:
         line_tokens = estimate_tokens(line)
         if current_tokens + line_tokens > effective_max_tokens:
             if current_chunk:
                 chunks.append("\n".join(current_chunk))
-            current_chunk = [line]
-            current_tokens = line_tokens
+            current_chunk, current_tokens = [line], line_tokens
         else:
             current_chunk.append(line)
             current_tokens += line_tokens
-
     if current_chunk:
         chunks.append("\n".join(current_chunk))
-
     return chunks
 
 def build_prompt_from_text(chunk: str) -> str:
@@ -113,10 +107,8 @@ Please analyze the above and provide:
 def init_agent():
     default_tool_path = os.path.abspath("data/new_tool.json")
     target_tool_path = os.path.join(tool_cache_dir, "new_tool.json")
-
     if not os.path.exists(target_tool_path):
         shutil.copy(default_tool_path, target_tool_path)
-
     agent = TxAgent(
         model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B",
         rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
@@ -141,7 +133,6 @@ def process_final_report(agent, file, chatbot_state: List[Dict[str, str]]) -> Tu
     try:
         messages.append({"role": "user", "content": f"Processing Excel file: {os.path.basename(file.name)}"})
         messages.append({"role": "assistant", "content": "⏳ Extracting and analyzing data..."})
-
         extracted_text = extract_text_from_excel(file.name)
         chunks = split_text_into_chunks(extracted_text)
         chunk_responses = [None] * len(chunks)
@@ -219,7 +210,6 @@ def process_final_report(agent, file, chatbot_state: List[Dict[str, str]]) -> Tu
                 except Exception as e:
                     messages.append({"role": "assistant", "content": f"❌ Error summarizing intermediate results: {str(e)}"})
                     return messages, report_path
-
             summary += f"\n\n### Chunk {i+1} Analysis\n{response}"
             current_summary_tokens += response_tokens
 
@@ -268,7 +258,7 @@ def process_final_report(agent, file, chatbot_state: List[Dict[str, str]]) -> Tu
 def create_ui(agent):
     with gr.Blocks(
         title="Patient History Chat",
-        css=\"\"\"
+        css="""
         .gradio-container {
             max-width: 900px !important;
             margin: auto;
@@ -303,24 +293,27 @@ def create_ui(agent):
             padding-left: 1.2em;
             margin: 0.4em 0;
         }
-        \"\"\"
+        """
     ) as demo:
-        gr.Markdown(\"\"\"\n<h2 style='color:#182848'>🏥 Patient History Analysis Tool</h2>\n<p style='color:#444;'>Upload an Excel file containing clinical data. The assistant will analyze it for patterns, inconsistencies, and recommendations.</p>\n\"\"\")
+        gr.Markdown("""
+        <h2 style='color:#182848'>🏥 Patient History Analysis Tool</h2>
+        <p style='color:#444;'>Upload an Excel file containing clinical data. The assistant will analyze it for patterns, inconsistencies, and recommendations.</p>
+        """)
 
         with gr.Row():
             with gr.Column(scale=3):
                 chatbot = gr.Chatbot(
-                    label=\"Clinical Assistant\",
+                    label="Clinical Assistant",
                     show_copy_button=True,
                     height=600,
-                    type=\"messages\",
-                    avatar_images=(None, \"https://i.imgur.com/6wX7Zb4.png\"),
+                    type="messages",
+                    avatar_images=(None, "https://i.imgur.com/6wX7Zb4.png"),
                     render_markdown=True
                 )
             with gr.Column(scale=1):
-                file_upload = gr.File(label=\"Upload Excel File\", file_types=[\".xlsx\"], height=100)
-                analyze_btn = gr.Button(\"🧠 Analyze Patient History\", variant=\"primary\", elem_classes=\"primary\")
-                report_output = gr.File(label=\"Download Report\", visible=False, interactive=False)
+                file_upload = gr.File(label="Upload Excel File", file_types=[".xlsx"], height=100)
+                analyze_btn = gr.Button("🧠 Analyze Patient History", variant="primary", elem_classes="primary")
+                report_output = gr.File(label="Download Report", visible=False, interactive=False)
 
         chatbot_state = gr.State(value=[])
 
@@ -328,24 +321,24 @@ def create_ui(agent):
             messages, report_path = process_final_report(agent, file, current_state)
             formatted_messages = []
             for msg in messages:
-                role = msg.get(\"role\")
-                content = msg.get(\"content\", \"\")
-                if role == \"assistant\":
-                    content = content.replace(\"- \", \"\\n- \")
-                    content = f\"<div class='chat-message-content'>{content}</div>\"
-                formatted_messages.append({\"role\": role, \"content\": content})
+                role = msg.get("role")
+                content = msg.get("content", "")
+                if role == "assistant":
+                    content = content.replace("- ", "\n- ")
+                    content = f"<div class='chat-message-content'>{content}</div>"
+                formatted_messages.append({"role": role, "content": content})
             report_update = gr.update(visible=report_path is not None, value=report_path)
             return formatted_messages, report_update, formatted_messages
 
-        analyze_btn.click(fn=update_ui, inputs=[file_upload, chatbot_state], outputs=[chatbot, report_output, chatbot_state], api_name=\"analyze\")
+        analyze_btn.click(fn=update_ui, inputs=[file_upload, chatbot_state], outputs=[chatbot, report_output, chatbot_state], api_name="analyze")
 
     return demo
 
-if __name__ == \"__main__\":
+if __name__ == "__main__":
     try:
         agent = init_agent()
         demo = create_ui(agent)
-        demo.launch(server_name=\"0.0.0.0\", server_port=7860, show_error=True, allowed_paths=[\"/data/hf_cache/reports\"], share=False)
+        demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True, allowed_paths=["/data/hf_cache/reports"], share=False)
     except Exception as e:
-        print(f\"Error: {str(e)}\")
+        print(f"Error: {str(e)}")
         sys.exit(1)