Update src/txagent/txagent.py

src/txagent/txagent.py

@@ -14,7 +14,6 @@ from .toolrag import ToolRAGModel
 import torch
 import logging
 from difflib import SequenceMatcher
-import asyncio
 import threading
 
 logger = logging.getLogger(__name__)
@@ -455,6 +454,140 @@ Patient Record Excerpt:
         logger.debug("Quick summary output: %s", output[:100])
         return output
 
+    def run_background_report(self, message: str, history: list, temperature: float,
+                             max_new_tokens: int, max_token: int, call_agent: bool,
+                             conversation: gr.State, max_round: int, seed: int,
+                             call_agent_level: int, report_path: str):
+        """Run detailed report generation in the background and save to file"""
+        logger.debug("Starting background report for message: %s", message[:100])
+        combined_response = ""
+        history_copy = history.copy()
+        
+        picked_tools_prompt, call_agent_level = self.initialize_tools_prompt(
+            call_agent, call_agent_level, message)
+        conversation = self.initialize_conversation(message, conversation, history_copy)
+        
+        next_round = True
+        current_round = 0
+        enable_summary = False
+        last_status = {}
+        token_overflow = False
+
+        if self.enable_checker:
+            checker = ReasoningTraceChecker(message, conversation, init_index=len(conversation))
+
+        try:
+            while next_round and current_round < max_round:
+                current_round += 1
+                last_outputs = []
+                if last_outputs:
+                    function_call_messages, picked_tools_prompt, special_tool_call = self.run_function_call(
+                        last_outputs, return_message=True, existing_tools_prompt=picked_tools_prompt,
+                        message_for_call_agent=message, call_agent=call_agent,
+                        call_agent_level=call_agent_level, temperature=temperature)
+                    
+                    if special_tool_call == 'Finish':
+                        next_round = False
+                        conversation.extend(function_call_messages)
+                        combined_response += function_call_messages[0]['content'] + "\n"
+                        break
+
+                    if (self.enable_summary or token_overflow) and not call_agent:
+                        enable_summary = True
+                    last_status = self.function_result_summary(
+                        conversation, status=last_status, enable_summary=enable_summary)
+
+                    if function_call_messages:
+                        conversation.extend(function_call_messages)
+                        combined_response += tool_result_format(function_call_messages) + "\n"
+                    else:
+                        next_round = False
+                        combined_response += ''.join(last_outputs).replace("</s>", "") + "\n"
+                        break
+
+                if self.enable_checker:
+                    good_status, wrong_info = checker.check_conversation()
+                    if not good_status:
+                        logger.warning("Checker error: %s", wrong_info)
+                        break
+
+                tools = picked_tools_prompt
+                last_outputs_str, token_overflow = self.llm_infer(
+                    messages=conversation, temperature=temperature, tools=tools,
+                    max_new_tokens=max_new_tokens, max_token=max_token, seed=seed, check_token_status=True)
+
+                if last_outputs_str is None:
+                    if self.force_finish:
+                        last_outputs_str = self.get_answer_based_on_unfinished_reasoning(
+                            conversation, temperature, max_new_tokens, max_token)
+                        combined_response += last_outputs_str + "\n"
+                        break
+                    combined_response += "Token limit exceeded.\n"
+                    break
+
+                combined_response += last_outputs_str + "\n"
+                last_outputs.append(last_outputs_str)
+
+            if next_round and self.force_finish:
+                last_outputs_str = self.get_answer_based_on_unfinished_reasoning(
+                    conversation, temperature, max_new_tokens, max_token)
+                combined_response += last_outputs_str + "\n"
+
+            # Save report
+            try:
+                with open(report_path, "w", encoding="utf-8") as f:
+                    f.write(combined_response)
+                logger.info("Detailed report saved to %s", report_path)
+            except Exception as e:
+                logger.error("Failed to save report: %s", e)
+
+        except Exception as e:
+            logger.error("Background report error: %s", e)
+            combined_response += f"Error: {e}\n"
+            with open(report_path, "w", encoding="utf-8") as f:
+                f.write(combined_response)
+
+        finally:
+            torch.cuda.empty_cache()
+            gc.collect()
+
+    def run_gradio_chat(self, message: str, history: list, temperature: float,
+                        max_new_tokens: int, max_token: int, call_agent: bool,
+                        conversation: gr.State, max_round: int = 3, seed: int = None,
+                        call_agent_level: int = 0, sub_agent_task: str = None,
+                        uploaded_files: list = None, report_path: str = None):
+        logger.debug("Chat started, message: %s", message[:100])
+        if not message or len(message.strip()) < 5:
+            yield "Please provide a valid message or upload files to analyze."
+            return
+
+        if message.startswith("[\U0001f9f0 Tool_RAG") or message.startswith("⚒️"):
+            return
+
+        clinical_keywords = ['medication', 'symptom', 'evaluation', 'diagnosis']
+        has_clinical_data = any(keyword in message.lower() for keyword in clinical_keywords)
+        call_agent = call_agent and not has_clinical_data
+
+        # Generate quick summary
+        quick_summary = self.run_quick_summary(
+            message, temperature=temperature, max_new_tokens=256, max_token=1024)
+        history.append(ChatMessage(role="assistant", content=f"**Quick Summary:**\n{quick_summary}"))
+        yield history
+
+        # Start background report generation
+        if report_path:
+            threading.Thread(
+                target=self.run_background_report,
+                args=(message, history, temperature, max_new_tokens, max_token, call_agent,
+                      conversation, max_round, seed, call_agent_level, report_path),
+                daemon=True
+            ).start()
+            history.append(ChatMessage(
+                role="assistant",
+                content="Generating detailed report in the background. Download will be available when ready."
+            ))
+            yield history
+
     def run_self_agent(self, message: str, temperature: float, max_new_tokens: int, max_token: int):
         logger.debug("Starting self agent")
         conversation = self.set_system_prompt([], self.self_prompt)
@@ -583,143 +716,4 @@ Summarize the function responses in one sentence with all necessary information.
                 setattr(self, key, value)
                 updated_attributes[key] = value
         logger.debug("Updated parameters: %s", updated_attributes)
-        return updated_attributes
-
-    async def run_background_report(self, message: str, history: list, temperature: float,
-                                   max_new_tokens: int, max_token: int, call_agent: bool,
-                                   conversation: gr.State, max_round: int, seed: int,
-                                   call_agent_level: int, report_path: str):
-        """Run detailed report generation in the background and save to file"""
-        logger.debug("Starting background report for message: %s", message[:100])
-        combined_response = ""
-        history_copy = history.copy()
-        
-        picked_tools_prompt, call_agent_level = self.initialize_tools_prompt(
-            call_agent, call_agent_level, message)
-        conversation = self.initialize_conversation(message, conversation, history_copy)
-        
-        next_round = True
-        current_round = 0
-        enable_summary = False
-        last_status = {}
-        token_overflow = False
-
-        if self.enable_checker:
-            checker = ReasoningTraceChecker(message, conversation, init_index=len(conversation))
-
-        try:
-            while next_round and current_round < max_round:
-                current_round += 1
-                last_outputs = []
-                if last_outputs:
-                    function_call_messages, picked_tools_prompt, special_tool_call, _ = yield from self.run_function_call_stream(
-                        last_outputs, return_message=True, existing_tools_prompt=picked_tools_prompt,
-                        message_for_call_agent=message, call_agent=call_agent,
-                        call_agent_level=call_agent_level, temperature=temperature,
-                        return_gradio_history=False)
-                    
-                    if special_tool_call == 'Finish':
-                        next_round = False
-                        conversation.extend(function_call_messages)
-                        combined_response += function_call_messages[0]['content'] + "\n"
-                        break
-
-                    if (self.enable_summary or token_overflow) and not call_agent:
-                        enable_summary = True
-                    last_status = self.function_result_summary(
-                        conversation, status=last_status, enable_summary=enable_summary)
-
-                    if function_call_messages:
-                        conversation.extend(function_call_messages)
-                        combined_response += tool_result_format(function_call_messages) + "\n"
-                    else:
-                        next_round = False
-                        combined_response += ''.join(last_outputs).replace("</s>", "") + "\n"
-                        break
-
-                if self.enable_checker:
-                    good_status, wrong_info = checker.check_conversation()
-                    if not good_status:
-                        logger.warning("Checker error: %s", wrong_info)
-                        break
-
-                tools = picked_tools_prompt
-                last_outputs_str, token_overflow = self.llm_infer(
-                    messages=conversation, temperature=temperature, tools=tools,
-                    max_new_tokens=max_new_tokens, max_token=max_token, seed=seed, check_token_status=True)
-
-                if last_outputs_str is None:
-                    if self.force_finish:
-                        last_outputs_str = self.get_answer_based_on_unfinished_reasoning(
-                            conversation, temperature, max_new_tokens, max_token)
-                        combined_response += last_outputs_str + "\n"
-                        break
-                    combined_response += "Token limit exceeded.\n"
-                    break
-
-                combined_response += last_outputs_str + "\n"
-                last_outputs.append(last_outputs_str)
-
-            if next_round and self.force_finish:
-                last_outputs_str = self.get_answer_based_on_unfinished_reasoning(
-                    conversation, temperature, max_new_tokens, max_token)
-                combined_response += last_outputs_str + "\n"
-
-            # Save report
-            try:
-                with open(report_path, "w", encoding="utf-8") as f:
-                    f.write(combined_response)
-                logger.info("Detailed report saved to %s", report_path)
-            except Exception as e:
-                logger.error("Failed to save report: %s", e)
-
-        except Exception as e:
-            logger.error("Background report error: %s", e)
-            combined_response += f"Error: {e}\n"
-            with open(report_path, "w", encoding="utf-8") as f:
-                f.write(combined_response)
-
-        finally:
-            torch.cuda.empty_cache()
-            gc.collect()
-
-    def run_gradio_chat(self, message: str, history: list, temperature: float,
-                        max_new_tokens: int, max_token: int, call_agent: bool,
-                        conversation: gr.State, max_round: int = 3, seed: int = None,
-                        call_agent_level: int = 0, sub_agent_task: str = None,
-                        uploaded_files: list = None, report_path: str = None):
-        logger.debug("Chat started, message: %s", message[:100])
-        if not message or len(message.strip()) < 5:
-            yield "Please provide a valid message or upload files to analyze."
-            return
-
-        if message.startswith("[\U0001f9f0 Tool_RAG") or message.startswith("⚒️"):
-            return
-
-        clinical_keywords = ['medication', 'symptom', 'evaluation', 'diagnosis']
-        has_clinical_data = any(keyword in message.lower() for keyword in clinical_keywords)
-        call_agent = call_agent and not has_clinical_data
-
-        # Generate quick summary
-        quick_summary = self.run_quick_summary(
-            message, temperature=temperature, max_new_tokens=256, max_token=1024)
-        history.append(ChatMessage(role="assistant", content=f"**Quick Summary:**\n{quick_summary}"))
-        yield history
-
-        # Start background report generation
-        if report_path:
-            loop = asyncio.get_event_loop()
-            threading.Thread(
-                target=lambda: loop.run_until_complete(
-                    self.run_background_report(
-                        message, history, temperature, max_new_tokens, max_token, call_agent,
-                        conversation, max_round, seed, call_agent_level, report_path
-                    )
-                ),
-                daemon=True
-            ).start()
-            history.append(ChatMessage(
-                role="assistant",
-                content="Generating detailed report in the background. Download will be available when ready."
-            ))
-            yield history
+        return updated_attributes